From 62391166891945952ca1e6ae7cd1e7510574cbc2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:18:59 +0000 Subject: [PATCH 1/4] Initial plan From c884c63eb40f8ae6c57ace7a4699db93f88cf137 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:27:49 +0000 Subject: [PATCH 2/4] Add SerpAPI AI Mode search wrangle Co-authored-by: ebhills <53243273+ebhills@users.noreply.github.com> --- README.md | 3 + docs/search-ai-mode.md | 182 ++++++++++ tests/test_search_ai_mode.py | 512 +++++++++++++++++++++++++++++ wrangles/__init__.py | 3 +- wrangles/clients/serp_api.py | 309 ++++++++++++++++- wrangles/format.py | 62 ++++ wrangles/recipe_wrangles/search.py | 204 ++++++++++++ wrangles/search.py | 106 ++++++ 8 files changed, 1378 insertions(+), 3 deletions(-) create mode 100644 docs/search-ai-mode.md create mode 100644 tests/test_search_ai_mode.py diff --git a/README.md b/README.md index d071f7c1b..08a6c3794 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,9 @@ Full documentation available at [wrangles.io](https://wrangles.io/python). +Documentation for the unreleased `search.ai_mode` wrangle is available in +[docs/search-ai-mode.md](docs/search-ai-mode.md). + ## Local development Supported local development uses Python 3.13. On Windows, create or refresh the diff --git a/docs/search-ai-mode.md b/docs/search-ai-mode.md new file mode 100644 index 000000000..ce10d62ca --- /dev/null +++ b/docs/search-ai-mode.md @@ -0,0 +1,182 @@ +# Search with Google AI Mode + +`search.ai_mode` uses SerpAPI's Google AI Mode API to search for sources and +return the synthesized, cited answer in one request. It is intended to replace +the common `search.find_links` followed by `search.retrieve_link_content` flow +when URL-by-URL retrieval is not required. + +## Industrial-product search + +The input query contains the product evidence assembled by the caller. For +example: + +```text +Manufacturer: SKF +Potential part codes: 6205-2RS, 6205 2RS +Description: deep groove ball bearing +``` + +By default, the wrangle asks AI Mode to find authoritative manufacturer, +product, supplier, and distributor pages; confirm exact manufacturer and part +identifiers; report important attributes and available pricing; cite sources; +distinguish facts from inference; and leave unknown values unknown. + +Set `prompt` to replace that instruction for a different research task. The +query cell is appended unchanged to either prompt, preserving manufacturer +names and part-code punctuation. + +## Recipe example + +```yaml +read: + - file: + name: products.csv + +wrangles: + - search.ai_mode: + queries: Product Search Query + id: ID + output: + - AI Mode Results + - AI Mode Text + n_results: 10 + country: us + language: en + +write: + - file: + name: researched-products.xlsx +``` + +`Product Search Query` may contain a scalar query or a list of queries in each +row. Multiple query columns are also supported when each has one corresponding +output column. A single query column may instead have exactly two output +columns: structured results followed by readable text. + +Blank, null, and `NaN` cells remain aligned as empty results and do not make a +provider request. The column named by `id` is copied to every source record as +`input_row_id`. + +## Direct Python API + +```python +import wrangles + +result = wrangles.search.ai_mode( + "Manufacturer: WESTFALIA\n" + "Potential part codes: DN65\n" + "Description: union nut", + n_results=10, + country="us", + language="en", +) +``` + +A scalar query returns one dictionary. A list returns an ordered list of +dictionaries. + +## Parameters + +| Parameter | Description | +| --- | --- | +| `queries` | Query column name(s) in recipes, or query value(s) in Python. | +| `id` | Recipe input row ID column. | +| `output` | Structured output column, or structured/text output pair. | +| `client` | `serpapi` (default and currently supported provider). | +| `api_key` | SerpAPI key; defaults to `SERPAPI_API_KEY`. | +| `prompt` | Optional replacement for the default product-research prompt. | +| `n_results` | Source limit after normalization and deduplication; minimum 1. | +| `threads` | Concurrent request count; minimum 1. | +| `country` | Friendly alias for SerpAPI `gl`; defaults to `us`. | +| `language` | Friendly alias for SerpAPI `hl`; defaults to `en`. | +| `location` | Human-readable search location. | +| `uule` | Encoded Google location; cannot be combined with `location`. | +| `device` | `desktop`, `tablet`, or `mobile`. | +| `no_cache` | Request a fresh result rather than a SerpAPI cached response. | +| `include_raw_response` | Add the provider response to each payload; defaults to `false`. | + +Other SerpAPI properties are not passed through by this wrangle. +`n_results` is applied locally and is not sent as the undocumented `num` +parameter. + +## Structured output + +Every query payload has the same top-level shape: + +```json +{ + "search_metadata": { + "query_index": 1, + "query": "Manufacturer: SKF ...", + "search_type": "ai_mode", + "search_id": "provider search id", + "status": "Success", + "search_date": null, + "response_time": null, + "json_endpoint": null, + "google_url": null, + "language": "en", + "country": "us", + "location": null + }, + "status": "Success", + "error": null, + "search_results": [ + { + "input_row_id": "row id", + "query_index": 1, + "google_rank": 1, + "result_type": "reference", + "title": "Source title", + "link": "example.com/product", + "source": "Example", + "snippet": "Supporting source snippet", + "pricing": { + "price": 12.5, + "currency": "USD", + "availability": "In stock", + "vendor": "Example" + } + } + ], + "extracted_content": { + "answer_markdown": "The synthesized answer with citations.", + "text_blocks": [] + } +} +``` + +SerpAPI `references`, source-bearing `quick_results`, `shopping_results`, and +`inline_products` become source records. `result_type` preserves their +provenance. Sources are deduplicated by cleaned link and title, ranked in +first-seen order, and then limited by `n_results`. Pricing is included only +when structured shopping data is available. + +`reconstructed_markdown` becomes `answer_markdown`; `text_blocks` is preserved. +The full provider payload is omitted unless `include_raw_response` is true. + +## Empty, partial, and error results + +- A successful response may have an answer but no sources or prices. +- Missing prices do not turn a successful response into a failure. +- Provider errors and request exceptions return `status: Failure`, a useful + `error`, empty `search_results`, and empty extracted content. +- One failed query does not shift or remove other query results. +- The readable output includes the query, status, error, answer, numbered + sources, snippets, and structured pricing that are present. + +## Cost and caching + +Each nonblank query may incur a SerpAPI Google AI Mode request and associated +provider charges. Review SerpAPI's current pricing and cache policy before +large batches. Cached responses can reduce repeated provider work; setting +`no_cache: true` requests a fresh response and may increase cost and latency. +Use `threads` to control concurrency and `n_results` to limit local output +size. + +## Unreleased release note + +Added `search.ai_mode` for one-request cited search and synthesis through +SerpAPI Google AI Mode, with stable normalized source/content output, readable +dual output, direct Python support, local source limiting, and opt-in raw +responses. diff --git a/tests/test_search_ai_mode.py b/tests/test_search_ai_mode.py new file mode 100644 index 000000000..1f01ca051 --- /dev/null +++ b/tests/test_search_ai_mode.py @@ -0,0 +1,512 @@ +import json + +import pandas as pd +import pytest +import yaml + +import wrangles +from wrangles import format as wrangles_format +from wrangles.clients.serp_api import SerpApiWranglesClient +from wrangles.recipe_wrangles import search as recipe_search + + +PRODUCT_QUERY = ( + "Manufacturer: SKF\n" + "Potential part codes: 6205-2RS, 6205 2RS\n" + "Description: deep groove ball bearing" +) + + +@pytest.fixture +def ai_mode_response(): + return { + "search_metadata": { + "id": "search-123", + "status": "Success", + "created_at": "2026-08-20 20:00:00 UTC", + "total_time_taken": 1.25, + "json_endpoint": "https://serpapi.com/searches/search-123.json", + "google_url": "https://www.google.com/search?q=SKF&utm_source=test", + }, + "search_parameters": { + "engine": "google_ai_mode", + "q": PRODUCT_QUERY, + "hl": "en", + "gl": "us", + "location_used": "Austin, Texas", + }, + "reconstructed_markdown": "SKF identifies **6205-2RS** as a sealed bearing.[1]", + "text_blocks": [ + { + "type": "paragraph", + "snippet": "SKF identifies 6205-2RS as a sealed bearing.", + "reference_indexes": [0], + } + ], + "references": [ + { + "title": "SKF 6205-2RS product page", + "link": "https://www.skf.com/products/6205-2RS?utm_source=google", + "source": "SKF", + "snippet": "Official product specifications.", + } + ], + "quick_results": [ + { + "title": "SKF 6205-2RS product page", + "link": "https://www.skf.com/products/6205-2RS", + "source": "SKF", + "snippet": "Official product specifications.", + }, + { + "title": "6205-2RS datasheet", + "link": "https://example.com/datasheet", + "source": "Example", + "snippet": " Bearing dimensions. | Bearing dimensions. ", + }, + ], + "shopping_results": [ + { + "title": "SKF 6205-2RS bearing", + "product_link": "https://supplier.example/skf-6205?gclid=tracking", + "source": "Bearing Supplier", + "snippet": "Available for immediate dispatch.", + "extracted_price": 12.5, + "currency": "USD", + "availability": "In stock", + } + ], + "inline_products": [ + { + "title": "SKF 6205 2RS", + "link": "https://distributor.example/6205", + "source": {"name": "Distributor", "link": "https://distributor.example"}, + "price": {"value": 14.25, "currency": "USD"}, + } + ], + } + + +class FakeSerpApiClient: + response = {} + requests = [] + + def __init__(self, api_key): + self.api_key = api_key + + def search(self, params): + self.requests.append(params) + if isinstance(self.response, Exception): + raise self.response + return self.response + + +def make_client(response): + FakeSerpApiClient.response = response + FakeSerpApiClient.requests = [] + client = SerpApiWranglesClient(api_key="test-key") + client.client_class = FakeSerpApiClient + return client + + +def normalized_payload(query, query_index=1): + return { + "search_metadata": { + "query_index": query_index, + "query": query, + "search_type": "ai_mode", + "search_id": "search-123", + "status": "Success", + "search_date": None, + "response_time": None, + "json_endpoint": None, + "google_url": None, + "language": "en", + "country": "us", + "location": None, + }, + "status": "Success", + "error": None, + "search_results": [ + { + "query_index": query_index, + "google_rank": 1, + "result_type": "reference", + "title": "Source title", + "link": "example.com/product", + "source": "Example", + "snippet": "Supporting source snippet", + "pricing": {}, + } + ], + "extracted_content": { + "answer_markdown": "A cited answer.", + "text_blocks": [], + }, + } + + +def test_client_maps_documented_ai_mode_response(ai_mode_response): + client = make_client(ai_mode_response) + + result = client.ai_mode_single( + PRODUCT_QUERY, + prompt="Research this exact industrial product.", + n_results=3, + query_index=2, + country="us", + language="en", + location="Austin, Texas", + device="desktop", + no_cache=True, + ) + + request = FakeSerpApiClient.requests[0] + assert request["engine"] == "google_ai_mode" + assert request["output"] == "json" + assert request["gl"] == "us" + assert request["hl"] == "en" + assert request["location"] == "Austin, Texas" + assert request["device"] == "desktop" + assert request["no_cache"] is True + assert "num" not in request + assert request["q"].startswith("Research this exact industrial product.") + assert "Find authoritative manufacturer" not in request["q"] + assert request["q"].endswith(PRODUCT_QUERY) + assert "6205-2RS, 6205 2RS" in request["q"] + + assert result["status"] == "Success" + assert result["error"] is None + assert result["search_metadata"]["query_index"] == 2 + assert result["search_metadata"]["query"] == PRODUCT_QUERY + assert result["search_metadata"]["search_type"] == "ai_mode" + assert result["extracted_content"] == { + "answer_markdown": "SKF identifies **6205-2RS** as a sealed bearing.[1]", + "text_blocks": ai_mode_response["text_blocks"], + } + assert len(result["search_results"]) == 3 + assert [item["result_type"] for item in result["search_results"]] == [ + "reference", + "quick_result", + "shopping_result", + ] + assert [item["google_rank"] for item in result["search_results"]] == [1, 2, 3] + assert result["search_results"][0]["link"] == "www.skf.com/products/6205-2RS" + assert result["search_results"][1]["snippet"] == "Bearing dimensions." + assert result["search_results"][2]["pricing"] == { + "price": 12.5, + "currency": "USD", + "availability": "In stock", + "vendor": "Bearing Supplier", + } + assert "raw_response" not in result + json.dumps(result) + + +def test_client_raw_response_and_partial_success(ai_mode_response): + ai_mode_response["references"] = [] + ai_mode_response["quick_results"] = [] + ai_mode_response["shopping_results"][0].pop("extracted_price") + ai_mode_response["shopping_results"][0].pop("currency") + client = make_client(ai_mode_response) + + result = client.ai_mode_single( + PRODUCT_QUERY, + n_results=10, + include_raw_response=True, + ) + + assert result["status"] == "Success" + assert result["search_results"][0]["pricing"] == { + "availability": "In stock", + "vendor": "Bearing Supplier", + } + assert result["raw_response"] == ai_mode_response + assert result["search_results"][1]["result_type"] == "inline_product" + assert result["search_results"][1]["pricing"] == { + "price": 14.25, + "currency": "USD", + "vendor": "Distributor", + } + json.dumps(result) + + +@pytest.mark.parametrize( + ("response", "message"), + [ + ({"search_metadata": {"status": "Error"}, "error": "API limit reached"}, "API limit reached"), + (RuntimeError("network unavailable"), "network unavailable"), + ], +) +def test_client_returns_stable_failure_payload(response, message): + client = make_client(response) + + result = client.ai_mode_single(PRODUCT_QUERY, query_index=1) + + assert result["status"] == "Failure" + assert result["error"] == message + assert result["search_results"] == [] + assert result["extracted_content"] == { + "answer_markdown": None, + "text_blocks": [], + } + + +def test_client_skips_blank_queries_without_provider_call(): + client = make_client({}) + + results = client.ai_mode_batch(["", None, float("nan")], threads=2) + + assert FakeSerpApiClient.requests == [] + assert len(results) == 3 + assert all(result["status"] == "Success" for result in results) + assert all(result["search_results"] == [] for result in results) + + +def test_client_batch_keeps_partial_failures_in_order(ai_mode_response): + class PartiallyFailingClient: + def __init__(self, api_key): + self.api_key = api_key + + def search(self, params): + if params["q"].endswith("bad query"): + raise RuntimeError("request failed") + return ai_mode_response + + client = make_client(ai_mode_response) + client.client_class = PartiallyFailingClient + + results = client.ai_mode_batch( + ["first query", "bad query", "last query"], + threads=3, + n_results=1, + ) + + assert [result["search_metadata"]["query"] for result in results] == [ + "first query", + "bad query", + "last query", + ] + assert [result["status"] for result in results] == [ + "Success", + "Failure", + "Success", + ] + assert results[1]["error"] == "request failed" + + +def test_direct_python_api_is_public_and_preserves_order(mocker): + calls = [] + + class FakeClient: + def ai_mode_batch(self, queries, **kwargs): + calls.append(kwargs) + values = queries if isinstance(queries, list) else [queries] + results = [ + normalized_payload(query, query_index=index) + for index, query in enumerate(values, start=1) + ] + return results if isinstance(queries, list) else results[0] + + factory = mocker.patch("wrangles.search._get_client", return_value=FakeClient()) + + scalar = wrangles.search.ai_mode(PRODUCT_QUERY, api_key="key") + multiple = wrangles.search.ai_mode(["first", "second"], api_key="key") + + assert isinstance(scalar, dict) + assert scalar["search_metadata"]["query"] == PRODUCT_QUERY + assert [item["search_metadata"]["query"] for item in multiple] == ["first", "second"] + assert calls[0]["prompt"].startswith( + "Find authoritative manufacturer, product, supplier, and distributor pages" + ) + assert factory.call_args.kwargs == { + "client_name": "serpapi", + "config": {"api_key": "key"}, + } + assert wrangles.search.SerpApiWranglesClient is SerpApiWranglesClient + + +def test_direct_python_api_returns_empty_without_creating_client(mocker): + factory = mocker.patch("wrangles.search._get_client") + + scalar = wrangles.search.ai_mode(None) + multiple = wrangles.search.ai_mode(["", float("nan")]) + + factory.assert_not_called() + assert scalar["search_results"] == [] + assert len(multiple) == 2 + assert all(result["status"] == "Success" for result in multiple) + + +def test_classic_find_links_public_api_remains_available(mocker): + class FakeClient: + def search_batch(self, queries, **kwargs): + return {"search_metadata": {"query": queries}, "search_results": []} + + mocker.patch("wrangles.search._get_client", return_value=FakeClient()) + + result = wrangles.search.find_links( + "classic search", + client_config={"api_key": "key"}, + n_results=3, + ) + + assert result == { + "search_metadata": {"query": "classic search"}, + "search_results": [], + } + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"n_results": 0}, "n_results must be at least 1"), + ({"threads": 0}, "threads must be at least 1"), + ({"device": "watch"}, "device must be one of"), + ({"location": "Austin", "uule": "encoded"}, "location and uule"), + ], +) +def test_direct_python_api_validates_parameters(kwargs, message): + with pytest.raises(ValueError, match=message): + wrangles.search.ai_mode(PRODUCT_QUERY, api_key="key", **kwargs) + + +def test_recipe_ai_mode_preserves_rows_queries_ids_and_readable_output(mocker): + calls = [] + + def fake_ai_mode(queries, **kwargs): + calls.append((queries, kwargs)) + return [ + normalized_payload(query, query_index=index) + for index, query in enumerate(queries, start=1) + ] + + mocker.patch.object(recipe_search._search_core, "ai_mode", side_effect=fake_ai_mode) + data = pd.DataFrame( + { + "query": [ + "Manufacturer: WESTFALIA\nPotential part code: DN65\nDescription: union nut", + ["general topic one", "general topic two"], + None, + float("nan"), + ], + "ID": ["westfalia", "general", "blank", "nan"], + } + ) + recipe = """ + wrangles: + - search.ai_mode: + queries: query + id: ID + output: + - AI Mode Results + - AI Mode Text + prompt: Find primary sources for this topic. + n_results: 4 + country: gb + language: en + device: mobile + no_cache: true + """ + + result = wrangles.recipe.run(recipe, dataframe=data) + + assert calls[0][0] == [ + "Manufacturer: WESTFALIA\nPotential part code: DN65\nDescription: union nut", + "general topic one", + "general topic two", + ] + assert calls[0][1]["prompt"] == "Find primary sources for this topic." + assert calls[0][1]["country"] == "gb" + assert calls[0][1]["device"] == "mobile" + assert result.loc[0, "AI Mode Results"][0]["search_results"][0]["input_row_id"] == "westfalia" + assert [ + payload["search_metadata"]["query_index"] + for payload in result.loc[1, "AI Mode Results"] + ] == [1, 2] + assert result.loc[2, "AI Mode Results"] == [] + assert result.loc[3, "AI Mode Results"] == [] + assert result.loc[2, "AI Mode Text"] == "" + assert "Query 1" in result.loc[0, "AI Mode Text"] + assert "A cited answer." in result.loc[0, "AI Mode Text"] + + +def test_recipe_ai_mode_supports_multiple_query_columns(mocker): + mocker.patch.object( + recipe_search._search_core, + "ai_mode", + side_effect=lambda queries, **kwargs: [ + normalized_payload(query, index) + for index, query in enumerate(queries, start=1) + ], + ) + data = pd.DataFrame( + { + "product_query": [PRODUCT_QUERY], + "general_query": ["History of ball bearings"], + "ID": [7], + } + ) + + result = recipe_search.ai_mode( + data, + queries=["product_query", "general_query"], + id="ID", + output=["product_results", "general_results"], + api_key="key", + ) + + assert result["product_results"][0][0]["search_metadata"]["query"] == PRODUCT_QUERY + assert result["general_results"][0][0]["search_metadata"]["query"] == "History of ball bearings" + + +def test_ai_mode_formatter_and_classic_formatter_contract(): + payload = normalized_payload(PRODUCT_QUERY) + payload["search_results"][0]["pricing"] = { + "price": 12.5, + "currency": "USD", + "availability": "In stock", + "vendor": "Example", + } + + text = wrangles_format.ai_mode_results_to_text([payload]) + classic_text = wrangles_format.raw_search_results_to_text([payload]) + + assert f"Query 1: {PRODUCT_QUERY}" in text + assert "Status: Success" in text + assert "A cited answer." in text + assert "Source 1" in text + assert "USD 12.5" in text + assert "## Query 1:" in classic_text + assert "Status: Success" not in classic_text + + failure = normalized_payload("failed query") + failure.update(status="Failure", error="API limit reached", search_results=[]) + failure_text = wrangles_format.ai_mode_results_to_text([failure]) + assert "Status: Failure" in failure_text + assert "Error: API limit reached" in failure_text + + +def test_ai_mode_recipe_schema_documents_supported_parameters_only(): + schema = yaml.safe_load(recipe_search.ai_mode.__doc__) + + assert schema["additionalProperties"] is False + assert schema["required"] == ["queries", "id", "output"] + assert set(schema["properties"]) == { + "queries", + "id", + "output", + "client", + "api_key", + "prompt", + "n_results", + "threads", + "country", + "language", + "location", + "uule", + "device", + "no_cache", + "include_raw_response", + } + assert schema["properties"]["device"]["enum"] == ["desktop", "tablet", "mobile"] diff --git a/wrangles/__init__.py b/wrangles/__init__.py index 97c3e2f63..21cf489f0 100644 --- a/wrangles/__init__.py +++ b/wrangles/__init__.py @@ -26,7 +26,7 @@ from . import ai_config from . import ai_definition from . import ai_cache -from .clients import serp_api as search +from . import search from . import data from .train import train @@ -35,4 +35,3 @@ from . import generate - diff --git a/wrangles/clients/serp_api.py b/wrangles/clients/serp_api.py index 962a9c9c0..8065d7797 100644 --- a/wrangles/clients/serp_api.py +++ b/wrangles/clients/serp_api.py @@ -1,4 +1,6 @@ import concurrent.futures as _futures +import json as _json +import math as _math import re from typing import Union as _Union @@ -94,6 +96,187 @@ def _extract_pricing_from_result(result: dict) -> dict: } +def _is_blank_query(query) -> bool: + if query is None: + return True + if isinstance(query, float) and _math.isnan(query): + return True + return str(query).strip().lower() in ("", "none", "nan", "nat") + + +def _json_safe(value): + return _json.loads(_json.dumps(value, default=str)) + + +def _ai_mode_payload( + query, + query_index: int | None, + *, + status: str = "Success", + error: str | None = None, + metadata: dict | None = None, + search_results: list | None = None, + answer_markdown=None, + text_blocks: list | None = None, +) -> dict: + metadata = metadata or {} + return { + "search_metadata": { + "query_index": query_index, + "query": None if query is None else str(query).strip(), + "search_type": "ai_mode", + "search_id": metadata.get("search_id"), + "status": status, + "search_date": metadata.get("search_date"), + "response_time": metadata.get("response_time"), + "json_endpoint": metadata.get("json_endpoint"), + "google_url": metadata.get("google_url"), + "language": metadata.get("language"), + "country": metadata.get("country"), + "location": metadata.get("location"), + }, + "status": status, + "error": error, + "search_results": search_results or [], + "extracted_content": { + "answer_markdown": answer_markdown, + "text_blocks": text_blocks or [], + }, + } + + +def _result_items(section) -> list[dict]: + if isinstance(section, list): + return [item for item in section if isinstance(item, dict)] + if not isinstance(section, dict): + return [] + for key in ("results", "items", "products"): + if isinstance(section.get(key), list): + return [item for item in section[key] if isinstance(item, dict)] + if any(key in section for key in ("link", "product_link", "title")): + return [section] + return [] + + +def _coerce_price(value): + if isinstance(value, (int, float)) and not isinstance(value, bool): + return value + if not isinstance(value, str): + return None + match = re.search(r"\d+(?:[.,]\d+)*", value) + if not match: + return None + number = match.group(0) + if "," in number and "." in number: + number = number.replace(",", "") + elif "," in number: + number = number.replace(",", ".") if re.search(r",\d{2}$", number) else number.replace(",", "") + try: + return float(number) + except ValueError: + return None + + +def _currency_from_price(value) -> str | None: + if not isinstance(value, str): + return None + for token, currency in ( + ("USD", "USD"), + ("CAD", "CAD"), + ("GBP", "GBP"), + ("EUR", "EUR"), + ("$", "USD"), + ("£", "GBP"), + ("€", "EUR"), + ): + haystack = value.upper() if token.isalpha() else value + if token in haystack: + return currency + return None + + +def _ai_mode_pricing(item: dict, source: str) -> dict: + raw_price = item.get("price") + if isinstance(raw_price, dict): + price = raw_price.get("value", raw_price.get("extracted_value")) + currency = raw_price.get("currency") + else: + price = item.get("extracted_price") + if price is None: + price = _coerce_price(raw_price) + currency = item.get("currency") or _currency_from_price(raw_price) + + availability = item.get("availability") or item.get("stock") + vendor = item.get("vendor") or item.get("merchant") or item.get("seller") or source + pricing = {} + if price is not None: + pricing["price"] = price + if currency: + pricing["currency"] = currency + if availability: + pricing["availability"] = availability + if vendor: + pricing["vendor"] = vendor + return pricing + + +def _ai_mode_result(item: dict, result_type: str, query_index: int | None) -> dict | None: + raw_source = item.get("source") + if isinstance(raw_source, dict): + source = raw_source.get("name") or raw_source.get("title") or "" + source_link = raw_source.get("link") or raw_source.get("url") + else: + source = raw_source or item.get("vendor") or item.get("merchant") or "" + source_link = None + + link = item.get("link") or item.get("product_link") or item.get("url") or source_link + if not link: + return None + + snippet = item.get("snippet") or item.get("description") or "" + result = { + "query_index": query_index, + "google_rank": 0, + "result_type": result_type, + "title": item.get("title") or item.get("name") or "", + "link": _web.clean_link(link), + "source": source, + "snippet": _web.clean_snippet(snippet), + "pricing": {}, + } + if result_type in ("shopping_result", "inline_product"): + result["pricing"] = _ai_mode_pricing(item, source) + return result + + +def _normalize_ai_mode_results(response: dict, query_index: int | None, n_results: int) -> list[dict]: + records = [] + seen = set() + sections = ( + ("references", "reference"), + ("quick_results", "quick_result"), + ("shopping_results", "shopping_result"), + ("inline_products", "inline_product"), + ) + for section_name, result_type in sections: + for item in _result_items(response.get(section_name)): + record = _ai_mode_result(item, result_type, query_index) + if record is None: + continue + key = ( + _web.normalize_site(record["link"]).lower().rstrip("/"), + record["title"].strip().lower(), + ) + if key in seen: + continue + seen.add(key) + records.append(record) + + for rank, record in enumerate(records[:n_results], start=1): + record["google_rank"] = rank + return records[:n_results] + + class SerpApiWranglesClient: def __init__(self, api_key: str = None): if not api_key or str(api_key).strip().lower() in ("", "none", "null"): @@ -220,4 +403,128 @@ def search_batch(self, input_data: _Union[str, list], n_results: int = 10, threa if input_was_scalar: return results[0] - return results \ No newline at end of file + return results + + def ai_mode_single( + self, + query, + prompt: str | None = None, + n_results: int = 10, + query_index: int | None = None, + country: str = "us", + language: str = "en", + location: str | None = None, + uule: str | None = None, + device: str = "desktop", + no_cache: bool = False, + include_raw_response: bool = False, + ) -> dict: + """Perform one Google AI Mode search and normalize the provider response.""" + if _is_blank_query(query): + return _ai_mode_payload( + query, + query_index, + metadata={ + "language": language, + "country": country, + "location": location, + }, + ) + + query_text = str(query).strip() + request_query = f"{prompt.strip()}\n\nQuery/product evidence:\n{query_text}" if prompt else query_text + params = { + "engine": "google_ai_mode", + "q": request_query, + "output": "json", + "gl": country, + "hl": language, + "device": device, + "no_cache": no_cache, + } + if location: + params["location"] = location + if uule: + params["uule"] = uule + + try: + client = self.client_class(api_key=self.api_key) + response = client.search(params) + if not isinstance(response, dict): + raise TypeError("SerpAPI returned a non-object response") + + meta_raw = response.get("search_metadata") or {} + search_params = response.get("search_parameters") or {} + provider_error = response.get("error") + provider_status = str(meta_raw.get("status") or "") + failed = bool(provider_error) or provider_status.lower() in ("error", "failed", "failure") + status = "Failure" if failed else "Success" + error = str(provider_error) if provider_error else ( + provider_status if failed else None + ) + metadata = { + "search_id": meta_raw.get("id"), + "search_date": meta_raw.get("created_at"), + "response_time": meta_raw.get("total_time_taken"), + "json_endpoint": meta_raw.get("json_endpoint"), + "google_url": _web.clean_link(meta_raw.get("google_url", "")) or None, + "language": search_params.get("hl", language), + "country": search_params.get("gl", country), + "location": search_params.get("location_used") or search_params.get("location") or location, + } + result = _ai_mode_payload( + query_text, + query_index, + status=status, + error=error, + metadata=metadata, + search_results=[] if failed else _normalize_ai_mode_results( + response, + query_index, + n_results, + ), + answer_markdown=None if failed else response.get("reconstructed_markdown"), + text_blocks=[] if failed else response.get("text_blocks"), + ) + if include_raw_response: + result["raw_response"] = _json_safe(response) + return result + except Exception as error: + return _ai_mode_payload( + query_text, + query_index, + status="Failure", + error=str(error), + metadata={ + "language": language, + "country": country, + "location": location, + }, + ) + + def ai_mode_batch( + self, + input_data: _Union[str, list], + prompt: str | None = None, + n_results: int = 10, + threads: int = 10, + **kwargs, + ) -> _Union[dict, list]: + """Perform ordered Google AI Mode searches in parallel.""" + input_was_scalar = not isinstance(input_data, list) + queries = [input_data] if input_was_scalar else input_data + indexed = list(enumerate(queries, start=1)) + + with _futures.ThreadPoolExecutor(max_workers=threads) as executor: + results = list(executor.map( + lambda item: self.ai_mode_single( + query=item[1], + prompt=prompt, + n_results=n_results, + query_index=item[0], + **kwargs, + ), + indexed, + )) + + return results[0] if input_was_scalar else results \ No newline at end of file diff --git a/wrangles/format.py b/wrangles/format.py index 1f4097280..121af3987 100644 --- a/wrangles/format.py +++ b/wrangles/format.py @@ -276,6 +276,68 @@ def raw_search_results_to_text(payloads: list) -> str: return "\n".join(lines).strip() +def ai_mode_results_to_text(payloads: list) -> str: + """Format normalized AI Mode payloads without changing classic search output.""" + if not payloads: + return "" + if not isinstance(payloads, list): + payloads = [payloads] + + blocks = [] + for index, payload in enumerate(payloads, start=1): + if not isinstance(payload, dict): + blocks.append(f"Query {index}\nStatus: Failure\nError: Invalid data") + continue + + metadata = payload.get("search_metadata") or {} + query_index = metadata.get("query_index") or index + query = metadata.get("query") or "" + status = payload.get("status") or metadata.get("status") or "Unknown" + lines = [f"## Query {query_index}: {query} ##", f"Status: {status}"] + + if payload.get("error"): + lines.append(f"Error: {payload['error']}") + + content = payload.get("extracted_content") or {} + answer = content.get("answer_markdown") if isinstance(content, dict) else content + if answer: + lines.extend(["", "Answer:", str(answer)]) + + results = payload.get("search_results") or [] + if results: + lines.extend(["", "Sources:"]) + for source_index, result in enumerate(results, start=1): + lines.append(f"# --- Source {source_index} --- #") + lines.append(f"Title: {result.get('title', '')}") + lines.append(f"Source: {result.get('source', '')}") + lines.append(f"Link: {result.get('link', '')}") + if result.get("snippet"): + snippet = textwrap.fill( + result["snippet"], + width=99, + subsequent_indent=" ", + ) + lines.append(f"Snippet: {snippet}") + + pricing = result.get("pricing") or {} + if pricing: + price = pricing.get("price") + currency = pricing.get("currency") + price_text = " ".join( + str(value) for value in (currency, price) if value is not None + ) or "Unknown" + pricing_parts = [price_text] + if pricing.get("availability"): + pricing_parts.append(str(pricing["availability"])) + if pricing.get("vendor"): + pricing_parts.append(f"via {pricing['vendor']}") + lines.append(f"Pricing: {' | '.join(pricing_parts)}") + + blocks.append("\n".join(lines)) + + return "\n\n========================================\n\n".join(blocks) + + def remove_duplicates(input_list: list, ignore_case: bool = False) -> list: """ Remove duplicates from a list. Preserves input order. diff --git a/wrangles/recipe_wrangles/search.py b/wrangles/recipe_wrangles/search.py index 11e8f08d7..074f122d2 100644 --- a/wrangles/recipe_wrangles/search.py +++ b/wrangles/recipe_wrangles/search.py @@ -171,6 +171,210 @@ def _to_query_list(v) -> list[str]: return df +def ai_mode( + df: _pd.DataFrame, + queries: str | list, + id: str, + output: str | list | None = None, + client: str = "serpapi", + api_key: str | None = None, + prompt: str | None = None, + n_results: int = 10, + threads: int = 10, + country: str = "us", + language: str = "en", + location: str | None = None, + uule: str | None = None, + device: str = "desktop", + no_cache: bool = False, + include_raw_response: bool = False, +) -> _pd.DataFrame: + """ + type: object + description: Search and synthesize cited content with SerpAPI Google AI Mode. + additionalProperties: false + required: + - queries + - id + - output + properties: + queries: + type: + - string + - array + description: Name or list of input columns containing query or product-evidence text. + id: + type: string + description: Name of the input row ID column copied to each source record. + output: + type: + - string + - array + description: Structured output column, or [structured_results, readable_text] for one query column. + client: + type: string + description: AI Mode search provider. + enum: + - serpapi + default: serpapi + api_key: + type: string + description: SerpAPI key. Defaults to the SERPAPI_API_KEY environment variable. + prompt: + type: string + description: Optional instruction replacing the default industrial-product research prompt. + n_results: + type: integer + minimum: 1 + description: Maximum normalized source records per query, applied after deduplication. + default: 10 + threads: + type: integer + minimum: 1 + description: Number of concurrent requests. + default: 10 + country: + type: string + description: Country code sent to SerpAPI as gl. + default: us + language: + type: string + description: Language code sent to SerpAPI as hl. + default: en + location: + type: string + description: Search location. Cannot be combined with uule. + uule: + type: string + description: Encoded Google location. Cannot be combined with location. + device: + type: string + description: Device type for the AI Mode request. + enum: + - desktop + - tablet + - mobile + default: desktop + no_cache: + type: boolean + description: Request a fresh SerpAPI response instead of cached results. + default: false + include_raw_response: + type: boolean + description: Include the JSON-safe provider response in each structured payload. + default: false + """ + if output is None: + output = queries + query_columns = queries if isinstance(queries, list) else [queries] + output_columns = output if isinstance(output, list) else [output] + is_dual_output = len(query_columns) == 1 and len(output_columns) == 2 + if not is_dual_output and len(query_columns) != len(output_columns): + raise ValueError( + "search.ai_mode must have an equal number of query and output columns, " + "OR 1 query column and 2 output columns [dicts, strings]." + ) + + def _is_blank(value) -> bool: + if value is None: + return True + if isinstance(value, str): + return not value.strip() + try: + missing = _pd.isna(value) + if not isinstance(missing, (list, tuple)): + return bool(missing) + except (TypeError, ValueError): + pass + return False + + def _to_query_list(value) -> list[str]: + if isinstance(value, (list, tuple)): + return [ + str(item).strip() + for item in value + if not _is_blank(item) + ] + return [] if _is_blank(value) else [str(value).strip()] + + row_ids = df[id].tolist() if id in df.columns else [None] * len(df) + for column_index, query_column in enumerate(query_columns): + structured_column = output_columns[0] if is_dual_output else output_columns[column_index] + row_query_lists = [_to_query_list(value) for value in df[query_column].tolist()] + flat_queries = [query for row_queries in row_query_lists for query in row_queries] + + if not flat_queries: + df[structured_column] = [[] for _ in row_query_lists] + if is_dual_output: + df[output_columns[1]] = ["" for _ in row_query_lists] + _logging.info(": Wrangling :: ai_mode summary :: 0 queries >> 0 results") + continue + + flat_responses = _search_core.ai_mode( + queries=flat_queries, + client=client, + api_key=api_key, + prompt=prompt, + n_results=n_results, + threads=threads, + country=country, + language=language, + location=location, + uule=uule, + device=device, + no_cache=no_cache, + include_raw_response=include_raw_response, + ) + if isinstance(flat_responses, dict): + flat_responses = [flat_responses] + + structured_cells = [] + text_cells = [] + position = 0 + total_results = 0 + for row_queries, current_id in zip(row_query_lists, row_ids): + query_count = len(row_queries) + if not query_count: + structured_cells.append([]) + text_cells.append("") + continue + + cell = flat_responses[position:position + query_count] + for query_index, response in enumerate(cell, start=1): + if not isinstance(response, dict): + continue + metadata = response.get("search_metadata") + if isinstance(metadata, dict): + metadata["query_index"] = query_index + source_records = [] + for source_record in response.get("search_results") or []: + if not isinstance(source_record, dict): + continue + updated_record = dict(source_record) + updated_record["input_row_id"] = current_id + updated_record["query_index"] = query_index + source_records.append(updated_record) + response["search_results"] = source_records + total_results += len(source_records) + + structured_cells.append(cell) + text_cells.append( + _format.ai_mode_results_to_text(cell) if is_dual_output else "" + ) + position += query_count + + df[structured_column] = structured_cells + if is_dual_output: + df[output_columns[1]] = text_cells + _logging.info( + f": Wrangling :: ai_mode summary :: {len(flat_queries)} queries " + f">> {total_results} results" + ) + + return df + + + def retrieve_link_content( df: _pd.DataFrame, diff --git a/wrangles/search.py b/wrangles/search.py index dedc359f4..a40aa81a4 100644 --- a/wrangles/search.py +++ b/wrangles/search.py @@ -1,7 +1,60 @@ import concurrent.futures as _futures +import math as _math # Import our client factory from .clients import get_client as _get_client +from .clients.serp_api import SerpApiWranglesClient + + +DEFAULT_AI_MODE_PROMPT = ( + "Find authoritative manufacturer, product, supplier, and distributor pages " + "for this industrial product. Confirm the manufacturer and exact part number " + "where possible. Summarize the product description, important specifications " + "and attributes, and available price, currency, vendor, availability, and " + "quantity basis. Prefer exact identifier evidence, distinguish confirmed facts " + "from inference, cite the supporting sources, and leave unknown values unknown " + "rather than guessing." +) + + +def _is_blank_ai_mode_query(query) -> bool: + if query is None: + return True + if isinstance(query, float) and _math.isnan(query): + return True + return str(query).strip().lower() in ("", "none", "nan", "nat") + + +def _empty_ai_mode_result( + query, + query_index: int, + country: str, + language: str, + location: str | None, +) -> dict: + return { + "search_metadata": { + "query_index": query_index, + "query": None if query is None else str(query).strip(), + "search_type": "ai_mode", + "search_id": None, + "status": "Success", + "search_date": None, + "response_time": None, + "json_endpoint": None, + "google_url": None, + "language": language, + "country": country, + "location": location, + }, + "status": "Success", + "error": None, + "search_results": [], + "extracted_content": { + "answer_markdown": None, + "text_blocks": [], + }, + } def find_links( @@ -27,6 +80,59 @@ def find_links( ) +def ai_mode( + queries: str | list, + client: str = "serpapi", + api_key: str | None = None, + prompt: str | None = None, + n_results: int = 10, + threads: int = 10, + country: str = "us", + language: str = "en", + location: str | None = None, + uule: str | None = None, + device: str = "desktop", + no_cache: bool = False, + include_raw_response: bool = False, +) -> dict | list: + """Search and synthesize cited content with SerpAPI Google AI Mode.""" + if not isinstance(n_results, int) or isinstance(n_results, bool) or n_results < 1: + raise ValueError("n_results must be at least 1") + if not isinstance(threads, int) or isinstance(threads, bool) or threads < 1: + raise ValueError("threads must be at least 1") + if location and uule: + raise ValueError("location and uule cannot both be provided") + if device not in ("desktop", "tablet", "mobile"): + raise ValueError("device must be one of: desktop, tablet, mobile") + + is_scalar = not isinstance(queries, list) + query_list = [queries] if is_scalar else queries + if all(_is_blank_ai_mode_query(query) for query in query_list): + empty_results = [ + _empty_ai_mode_result(query, index, country, language, location) + for index, query in enumerate(query_list, start=1) + ] + return empty_results[0] if is_scalar else empty_results + + search_client = _get_client( + client_name=client, + config={"api_key": api_key}, + ) + return search_client.ai_mode_batch( + queries, + prompt=DEFAULT_AI_MODE_PROMPT if prompt is None else prompt, + n_results=n_results, + threads=threads, + country=country, + language=language, + location=location, + uule=uule, + device=device, + no_cache=no_cache, + include_raw_response=include_raw_response, + ) + + def retrieve_link_content( urls: str | list, client: str = "google_url_context", From d39212aa38838496d1afc536e22e135a13624b0d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:42:02 +0000 Subject: [PATCH 3/4] Fix AI Mode metadata and deduplication Co-authored-by: ebhills <53243273+ebhills@users.noreply.github.com> --- .../wrangles}/test_search_ai_mode.py | 41 ++++++++++++++++++- wrangles/clients/serp_api.py | 38 +++++++++++++---- 2 files changed, 70 insertions(+), 9 deletions(-) rename tests/{ => recipes/wrangles}/test_search_ai_mode.py (93%) diff --git a/tests/test_search_ai_mode.py b/tests/recipes/wrangles/test_search_ai_mode.py similarity index 93% rename from tests/test_search_ai_mode.py rename to tests/recipes/wrangles/test_search_ai_mode.py index 1f01ca051..5643c67c0 100644 --- a/tests/test_search_ai_mode.py +++ b/tests/recipes/wrangles/test_search_ai_mode.py @@ -26,7 +26,7 @@ def ai_mode_response(): "created_at": "2026-08-20 20:00:00 UTC", "total_time_taken": 1.25, "json_endpoint": "https://serpapi.com/searches/search-123.json", - "google_url": "https://www.google.com/search?q=SKF&utm_source=test", + "google_ai_mode_url": "https://www.google.com/search?q=SKF&utm_source=test", }, "search_parameters": { "engine": "google_ai_mode", @@ -54,7 +54,7 @@ def ai_mode_response(): "quick_results": [ { "title": "SKF 6205-2RS product page", - "link": "https://www.skf.com/products/6205-2RS", + "link": "https://skf.com/products/6205-2RS", "source": "SKF", "snippet": "Official product specifications.", }, @@ -180,6 +180,7 @@ def test_client_maps_documented_ai_mode_response(ai_mode_response): assert result["search_metadata"]["query_index"] == 2 assert result["search_metadata"]["query"] == PRODUCT_QUERY assert result["search_metadata"]["search_type"] == "ai_mode" + assert result["search_metadata"]["google_url"] == "www.google.com/search?q=SKF" assert result["extracted_content"] == { "answer_markdown": "SKF identifies **6205-2RS** as a sealed bearing.[1]", "text_blocks": ai_mode_response["text_blocks"], @@ -231,6 +232,42 @@ def test_client_raw_response_and_partial_success(ai_mode_response): json.dumps(result) +@pytest.mark.parametrize( + ("formatted_price", "expected_currency"), + [ + ("CA$12.50", "CAD"), + ("A$12.50", "AUD"), + ("US$12.50", "USD"), + ("$12.50", None), + ], +) +def test_client_does_not_guess_ambiguous_price_currency( + formatted_price, + expected_currency, +): + response = { + "search_metadata": {"status": "Success"}, + "shopping_results": [ + { + "title": "Industrial product", + "link": "https://supplier.example/product", + "source": "Supplier", + "price": formatted_price, + } + ], + } + client = make_client(response) + + result = client.ai_mode_single("industrial product") + + pricing = result["search_results"][0]["pricing"] + assert pricing["price"] == 12.5 + if expected_currency: + assert pricing["currency"] == expected_currency + else: + assert "currency" not in pricing + + @pytest.mark.parametrize( ("response", "message"), [ diff --git a/wrangles/clients/serp_api.py b/wrangles/clients/serp_api.py index 8065d7797..b060e8c01 100644 --- a/wrangles/clients/serp_api.py +++ b/wrangles/clients/serp_api.py @@ -180,17 +180,24 @@ def _coerce_price(value): def _currency_from_price(value) -> str | None: if not isinstance(value, str): return None + value_upper = value.upper() for token, currency in ( + ("CA$", "CAD"), + ("C$", "CAD"), + ("A$", "AUD"), + ("AU$", "AUD"), + ("NZ$", "NZD"), + ("US$", "USD"), ("USD", "USD"), ("CAD", "CAD"), + ("AUD", "AUD"), + ("NZD", "NZD"), ("GBP", "GBP"), ("EUR", "EUR"), - ("$", "USD"), ("£", "GBP"), ("€", "EUR"), ): - haystack = value.upper() if token.isalpha() else value - if token in haystack: + if token in value_upper: return currency return None @@ -263,8 +270,11 @@ def _normalize_ai_mode_results(response: dict, query_index: int | None, n_result record = _ai_mode_result(item, result_type, query_index) if record is None: continue + dedupe_link = record["link"] + if "://" not in dedupe_link: + dedupe_link = f"https://{dedupe_link}" key = ( - _web.normalize_site(record["link"]).lower().rstrip("/"), + _web.normalize_site(dedupe_link).lower().rstrip("/"), record["title"].strip().lower(), ) if key in seen: @@ -467,7 +477,9 @@ def ai_mode_single( "search_date": meta_raw.get("created_at"), "response_time": meta_raw.get("total_time_taken"), "json_endpoint": meta_raw.get("json_endpoint"), - "google_url": _web.clean_link(meta_raw.get("google_url", "")) or None, + "google_url": _web.clean_link( + meta_raw.get("google_ai_mode_url") or meta_raw.get("google_url", "") + ) or None, "language": search_params.get("hl", language), "country": search_params.get("gl", country), "location": search_params.get("location_used") or search_params.get("location") or location, @@ -508,7 +520,13 @@ def ai_mode_batch( prompt: str | None = None, n_results: int = 10, threads: int = 10, - **kwargs, + country: str = "us", + language: str = "en", + location: str | None = None, + uule: str | None = None, + device: str = "desktop", + no_cache: bool = False, + include_raw_response: bool = False, ) -> _Union[dict, list]: """Perform ordered Google AI Mode searches in parallel.""" input_was_scalar = not isinstance(input_data, list) @@ -522,7 +540,13 @@ def ai_mode_batch( prompt=prompt, n_results=n_results, query_index=item[0], - **kwargs, + country=country, + language=language, + location=location, + uule=uule, + device=device, + no_cache=no_cache, + include_raw_response=include_raw_response, ), indexed, )) From a084f13b3a1864c5c12d8329fd2f5e83f1fa52bb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 21 Aug 2026 16:13:13 +0000 Subject: [PATCH 4/4] Address AI Mode review feedback Co-authored-by: thomasstvr <65630492+thomasstvr@users.noreply.github.com> --- docs/search-ai-mode.md | 22 +- tests/recipes/wrangles/test_search.py | 102 ++++ tests/recipes/wrangles/test_search_ai_mode.py | 549 ------------------ tests/test_wrangles.py | 20 +- wrangles/clients/serp_api.py | 28 +- wrangles/recipe_wrangles/search.py | 24 +- wrangles/search.py | 14 +- 7 files changed, 142 insertions(+), 617 deletions(-) delete mode 100644 tests/recipes/wrangles/test_search_ai_mode.py diff --git a/docs/search-ai-mode.md b/docs/search-ai-mode.md index ce10d62ca..fb81bdeb2 100644 --- a/docs/search-ai-mode.md +++ b/docs/search-ai-mode.md @@ -39,7 +39,6 @@ wrangles: output: - AI Mode Results - AI Mode Text - n_results: 10 country: us language: en @@ -66,7 +65,6 @@ result = wrangles.search.ai_mode( "Manufacturer: WESTFALIA\n" "Potential part codes: DN65\n" "Description: union nut", - n_results=10, country="us", language="en", ) @@ -85,19 +83,15 @@ dictionaries. | `client` | `serpapi` (default and currently supported provider). | | `api_key` | SerpAPI key; defaults to `SERPAPI_API_KEY`. | | `prompt` | Optional replacement for the default product-research prompt. | -| `n_results` | Source limit after normalization and deduplication; minimum 1. | | `threads` | Concurrent request count; minimum 1. | | `country` | Friendly alias for SerpAPI `gl`; defaults to `us`. | | `language` | Friendly alias for SerpAPI `hl`; defaults to `en`. | | `location` | Human-readable search location. | -| `uule` | Encoded Google location; cannot be combined with `location`. | -| `device` | `desktop`, `tablet`, or `mobile`. | | `no_cache` | Request a fresh result rather than a SerpAPI cached response. | | `include_raw_response` | Add the provider response to each payload; defaults to `false`. | -Other SerpAPI properties are not passed through by this wrangle. -`n_results` is applied locally and is not sent as the undocumented `num` -parameter. +Other SerpAPI properties may be passed as keyword arguments through the direct +Python API. AI Mode requests always use the supported desktop device. ## Structured output @@ -148,9 +142,9 @@ Every query payload has the same top-level shape: SerpAPI `references`, source-bearing `quick_results`, `shopping_results`, and `inline_products` become source records. `result_type` preserves their -provenance. Sources are deduplicated by cleaned link and title, ranked in -first-seen order, and then limited by `n_results`. Pricing is included only -when structured shopping data is available. +provenance. Sources are deduplicated by cleaned link and title and ranked in +first-seen order. Pricing is included only when structured shopping data is +available. `reconstructed_markdown` becomes `answer_markdown`; `text_blocks` is preserved. The full provider payload is omitted unless `include_raw_response` is true. @@ -171,12 +165,10 @@ Each nonblank query may incur a SerpAPI Google AI Mode request and associated provider charges. Review SerpAPI's current pricing and cache policy before large batches. Cached responses can reduce repeated provider work; setting `no_cache: true` requests a fresh response and may increase cost and latency. -Use `threads` to control concurrency and `n_results` to limit local output -size. +Use `threads` to control concurrency. ## Unreleased release note Added `search.ai_mode` for one-request cited search and synthesis through SerpAPI Google AI Mode, with stable normalized source/content output, readable -dual output, direct Python support, local source limiting, and opt-in raw -responses. +dual output, direct Python support, and opt-in raw responses. diff --git a/tests/recipes/wrangles/test_search.py b/tests/recipes/wrangles/test_search.py index dc2c69883..bf270acb1 100644 --- a/tests/recipes/wrangles/test_search.py +++ b/tests/recipes/wrangles/test_search.py @@ -398,6 +398,108 @@ def test_numeric_input_column(self): assert all(isinstance(row['results'][0]['search_results'], list) for _, row in df.iterrows()) +class TestAiMode: + query = "SKF 6205-2RS deep groove ball bearing specifications" + + def test_search_single_query(self): + data = pd.DataFrame({ + "query": [self.query], + "ID": ["bearing"], + }) + recipe = """ + wrangles: + - search.ai_mode: + queries: query + id: ID + output: results + api_key: ${SERPAPI_API_KEY} + country: us + language: en + location: Austin, Texas, United States + """ + + df = wrangles.recipe.run(recipe, dataframe=data) + + result = df.iloc[0]["results"][0] + assert result["status"] == "Success", result["error"] + assert result["search_metadata"]["query"] == self.query + assert result["search_metadata"]["location"] + assert isinstance(result["search_results"], list) + assert ( + result["extracted_content"]["answer_markdown"] + or result["extracted_content"]["text_blocks"] + ) + + def test_search_multiple_queries(self): + queries = [ + self.query, + "SKF 6205-2RS bearing dimensions", + ] + data = pd.DataFrame({ + "query": [queries], + "ID": ["bearing"], + }) + recipe = """ + wrangles: + - search.ai_mode: + queries: query + id: ID + output: results + api_key: ${SERPAPI_API_KEY} + """ + + df = wrangles.recipe.run(recipe, dataframe=data) + + results = df.iloc[0]["results"] + assert len(results) == 2 + assert [result["search_metadata"]["query"] for result in results] == queries + assert all(result["status"] == "Success" for result in results) + assert all( + source["input_row_id"] == "bearing" + for result in results + for source in result["search_results"] + ) + + def test_search_structured_and_readable_outputs(self): + data = pd.DataFrame({ + "query": [self.query], + "ID": ["bearing"], + }) + recipe = """ + wrangles: + - search.ai_mode: + queries: query + id: ID + output: + - results + - result_text + api_key: ${SERPAPI_API_KEY} + """ + + df = wrangles.recipe.run(recipe, dataframe=data) + + assert df.iloc[0]["results"][0]["status"] == "Success" + assert "Query 1:" in df.iloc[0]["result_text"] + assert self.query in df.iloc[0]["result_text"] + + def test_search_empty_input(self): + data = pd.DataFrame({ + "query": ["", None], + "ID": [1, 2], + }) + recipe = """ + wrangles: + - search.ai_mode: + queries: query + id: ID + output: results + """ + + df = wrangles.recipe.run(recipe, dataframe=data) + + assert df["results"].tolist() == [[], []] + + class TestRetrieveLinkContent: """ Test the functionality of the retrieve_link_content wrangle diff --git a/tests/recipes/wrangles/test_search_ai_mode.py b/tests/recipes/wrangles/test_search_ai_mode.py deleted file mode 100644 index 5643c67c0..000000000 --- a/tests/recipes/wrangles/test_search_ai_mode.py +++ /dev/null @@ -1,549 +0,0 @@ -import json - -import pandas as pd -import pytest -import yaml - -import wrangles -from wrangles import format as wrangles_format -from wrangles.clients.serp_api import SerpApiWranglesClient -from wrangles.recipe_wrangles import search as recipe_search - - -PRODUCT_QUERY = ( - "Manufacturer: SKF\n" - "Potential part codes: 6205-2RS, 6205 2RS\n" - "Description: deep groove ball bearing" -) - - -@pytest.fixture -def ai_mode_response(): - return { - "search_metadata": { - "id": "search-123", - "status": "Success", - "created_at": "2026-08-20 20:00:00 UTC", - "total_time_taken": 1.25, - "json_endpoint": "https://serpapi.com/searches/search-123.json", - "google_ai_mode_url": "https://www.google.com/search?q=SKF&utm_source=test", - }, - "search_parameters": { - "engine": "google_ai_mode", - "q": PRODUCT_QUERY, - "hl": "en", - "gl": "us", - "location_used": "Austin, Texas", - }, - "reconstructed_markdown": "SKF identifies **6205-2RS** as a sealed bearing.[1]", - "text_blocks": [ - { - "type": "paragraph", - "snippet": "SKF identifies 6205-2RS as a sealed bearing.", - "reference_indexes": [0], - } - ], - "references": [ - { - "title": "SKF 6205-2RS product page", - "link": "https://www.skf.com/products/6205-2RS?utm_source=google", - "source": "SKF", - "snippet": "Official product specifications.", - } - ], - "quick_results": [ - { - "title": "SKF 6205-2RS product page", - "link": "https://skf.com/products/6205-2RS", - "source": "SKF", - "snippet": "Official product specifications.", - }, - { - "title": "6205-2RS datasheet", - "link": "https://example.com/datasheet", - "source": "Example", - "snippet": " Bearing dimensions. | Bearing dimensions. ", - }, - ], - "shopping_results": [ - { - "title": "SKF 6205-2RS bearing", - "product_link": "https://supplier.example/skf-6205?gclid=tracking", - "source": "Bearing Supplier", - "snippet": "Available for immediate dispatch.", - "extracted_price": 12.5, - "currency": "USD", - "availability": "In stock", - } - ], - "inline_products": [ - { - "title": "SKF 6205 2RS", - "link": "https://distributor.example/6205", - "source": {"name": "Distributor", "link": "https://distributor.example"}, - "price": {"value": 14.25, "currency": "USD"}, - } - ], - } - - -class FakeSerpApiClient: - response = {} - requests = [] - - def __init__(self, api_key): - self.api_key = api_key - - def search(self, params): - self.requests.append(params) - if isinstance(self.response, Exception): - raise self.response - return self.response - - -def make_client(response): - FakeSerpApiClient.response = response - FakeSerpApiClient.requests = [] - client = SerpApiWranglesClient(api_key="test-key") - client.client_class = FakeSerpApiClient - return client - - -def normalized_payload(query, query_index=1): - return { - "search_metadata": { - "query_index": query_index, - "query": query, - "search_type": "ai_mode", - "search_id": "search-123", - "status": "Success", - "search_date": None, - "response_time": None, - "json_endpoint": None, - "google_url": None, - "language": "en", - "country": "us", - "location": None, - }, - "status": "Success", - "error": None, - "search_results": [ - { - "query_index": query_index, - "google_rank": 1, - "result_type": "reference", - "title": "Source title", - "link": "example.com/product", - "source": "Example", - "snippet": "Supporting source snippet", - "pricing": {}, - } - ], - "extracted_content": { - "answer_markdown": "A cited answer.", - "text_blocks": [], - }, - } - - -def test_client_maps_documented_ai_mode_response(ai_mode_response): - client = make_client(ai_mode_response) - - result = client.ai_mode_single( - PRODUCT_QUERY, - prompt="Research this exact industrial product.", - n_results=3, - query_index=2, - country="us", - language="en", - location="Austin, Texas", - device="desktop", - no_cache=True, - ) - - request = FakeSerpApiClient.requests[0] - assert request["engine"] == "google_ai_mode" - assert request["output"] == "json" - assert request["gl"] == "us" - assert request["hl"] == "en" - assert request["location"] == "Austin, Texas" - assert request["device"] == "desktop" - assert request["no_cache"] is True - assert "num" not in request - assert request["q"].startswith("Research this exact industrial product.") - assert "Find authoritative manufacturer" not in request["q"] - assert request["q"].endswith(PRODUCT_QUERY) - assert "6205-2RS, 6205 2RS" in request["q"] - - assert result["status"] == "Success" - assert result["error"] is None - assert result["search_metadata"]["query_index"] == 2 - assert result["search_metadata"]["query"] == PRODUCT_QUERY - assert result["search_metadata"]["search_type"] == "ai_mode" - assert result["search_metadata"]["google_url"] == "www.google.com/search?q=SKF" - assert result["extracted_content"] == { - "answer_markdown": "SKF identifies **6205-2RS** as a sealed bearing.[1]", - "text_blocks": ai_mode_response["text_blocks"], - } - assert len(result["search_results"]) == 3 - assert [item["result_type"] for item in result["search_results"]] == [ - "reference", - "quick_result", - "shopping_result", - ] - assert [item["google_rank"] for item in result["search_results"]] == [1, 2, 3] - assert result["search_results"][0]["link"] == "www.skf.com/products/6205-2RS" - assert result["search_results"][1]["snippet"] == "Bearing dimensions." - assert result["search_results"][2]["pricing"] == { - "price": 12.5, - "currency": "USD", - "availability": "In stock", - "vendor": "Bearing Supplier", - } - assert "raw_response" not in result - json.dumps(result) - - -def test_client_raw_response_and_partial_success(ai_mode_response): - ai_mode_response["references"] = [] - ai_mode_response["quick_results"] = [] - ai_mode_response["shopping_results"][0].pop("extracted_price") - ai_mode_response["shopping_results"][0].pop("currency") - client = make_client(ai_mode_response) - - result = client.ai_mode_single( - PRODUCT_QUERY, - n_results=10, - include_raw_response=True, - ) - - assert result["status"] == "Success" - assert result["search_results"][0]["pricing"] == { - "availability": "In stock", - "vendor": "Bearing Supplier", - } - assert result["raw_response"] == ai_mode_response - assert result["search_results"][1]["result_type"] == "inline_product" - assert result["search_results"][1]["pricing"] == { - "price": 14.25, - "currency": "USD", - "vendor": "Distributor", - } - json.dumps(result) - - -@pytest.mark.parametrize( - ("formatted_price", "expected_currency"), - [ - ("CA$12.50", "CAD"), - ("A$12.50", "AUD"), - ("US$12.50", "USD"), - ("$12.50", None), - ], -) -def test_client_does_not_guess_ambiguous_price_currency( - formatted_price, - expected_currency, -): - response = { - "search_metadata": {"status": "Success"}, - "shopping_results": [ - { - "title": "Industrial product", - "link": "https://supplier.example/product", - "source": "Supplier", - "price": formatted_price, - } - ], - } - client = make_client(response) - - result = client.ai_mode_single("industrial product") - - pricing = result["search_results"][0]["pricing"] - assert pricing["price"] == 12.5 - if expected_currency: - assert pricing["currency"] == expected_currency - else: - assert "currency" not in pricing - - -@pytest.mark.parametrize( - ("response", "message"), - [ - ({"search_metadata": {"status": "Error"}, "error": "API limit reached"}, "API limit reached"), - (RuntimeError("network unavailable"), "network unavailable"), - ], -) -def test_client_returns_stable_failure_payload(response, message): - client = make_client(response) - - result = client.ai_mode_single(PRODUCT_QUERY, query_index=1) - - assert result["status"] == "Failure" - assert result["error"] == message - assert result["search_results"] == [] - assert result["extracted_content"] == { - "answer_markdown": None, - "text_blocks": [], - } - - -def test_client_skips_blank_queries_without_provider_call(): - client = make_client({}) - - results = client.ai_mode_batch(["", None, float("nan")], threads=2) - - assert FakeSerpApiClient.requests == [] - assert len(results) == 3 - assert all(result["status"] == "Success" for result in results) - assert all(result["search_results"] == [] for result in results) - - -def test_client_batch_keeps_partial_failures_in_order(ai_mode_response): - class PartiallyFailingClient: - def __init__(self, api_key): - self.api_key = api_key - - def search(self, params): - if params["q"].endswith("bad query"): - raise RuntimeError("request failed") - return ai_mode_response - - client = make_client(ai_mode_response) - client.client_class = PartiallyFailingClient - - results = client.ai_mode_batch( - ["first query", "bad query", "last query"], - threads=3, - n_results=1, - ) - - assert [result["search_metadata"]["query"] for result in results] == [ - "first query", - "bad query", - "last query", - ] - assert [result["status"] for result in results] == [ - "Success", - "Failure", - "Success", - ] - assert results[1]["error"] == "request failed" - - -def test_direct_python_api_is_public_and_preserves_order(mocker): - calls = [] - - class FakeClient: - def ai_mode_batch(self, queries, **kwargs): - calls.append(kwargs) - values = queries if isinstance(queries, list) else [queries] - results = [ - normalized_payload(query, query_index=index) - for index, query in enumerate(values, start=1) - ] - return results if isinstance(queries, list) else results[0] - - factory = mocker.patch("wrangles.search._get_client", return_value=FakeClient()) - - scalar = wrangles.search.ai_mode(PRODUCT_QUERY, api_key="key") - multiple = wrangles.search.ai_mode(["first", "second"], api_key="key") - - assert isinstance(scalar, dict) - assert scalar["search_metadata"]["query"] == PRODUCT_QUERY - assert [item["search_metadata"]["query"] for item in multiple] == ["first", "second"] - assert calls[0]["prompt"].startswith( - "Find authoritative manufacturer, product, supplier, and distributor pages" - ) - assert factory.call_args.kwargs == { - "client_name": "serpapi", - "config": {"api_key": "key"}, - } - assert wrangles.search.SerpApiWranglesClient is SerpApiWranglesClient - - -def test_direct_python_api_returns_empty_without_creating_client(mocker): - factory = mocker.patch("wrangles.search._get_client") - - scalar = wrangles.search.ai_mode(None) - multiple = wrangles.search.ai_mode(["", float("nan")]) - - factory.assert_not_called() - assert scalar["search_results"] == [] - assert len(multiple) == 2 - assert all(result["status"] == "Success" for result in multiple) - - -def test_classic_find_links_public_api_remains_available(mocker): - class FakeClient: - def search_batch(self, queries, **kwargs): - return {"search_metadata": {"query": queries}, "search_results": []} - - mocker.patch("wrangles.search._get_client", return_value=FakeClient()) - - result = wrangles.search.find_links( - "classic search", - client_config={"api_key": "key"}, - n_results=3, - ) - - assert result == { - "search_metadata": {"query": "classic search"}, - "search_results": [], - } - - -@pytest.mark.parametrize( - ("kwargs", "message"), - [ - ({"n_results": 0}, "n_results must be at least 1"), - ({"threads": 0}, "threads must be at least 1"), - ({"device": "watch"}, "device must be one of"), - ({"location": "Austin", "uule": "encoded"}, "location and uule"), - ], -) -def test_direct_python_api_validates_parameters(kwargs, message): - with pytest.raises(ValueError, match=message): - wrangles.search.ai_mode(PRODUCT_QUERY, api_key="key", **kwargs) - - -def test_recipe_ai_mode_preserves_rows_queries_ids_and_readable_output(mocker): - calls = [] - - def fake_ai_mode(queries, **kwargs): - calls.append((queries, kwargs)) - return [ - normalized_payload(query, query_index=index) - for index, query in enumerate(queries, start=1) - ] - - mocker.patch.object(recipe_search._search_core, "ai_mode", side_effect=fake_ai_mode) - data = pd.DataFrame( - { - "query": [ - "Manufacturer: WESTFALIA\nPotential part code: DN65\nDescription: union nut", - ["general topic one", "general topic two"], - None, - float("nan"), - ], - "ID": ["westfalia", "general", "blank", "nan"], - } - ) - recipe = """ - wrangles: - - search.ai_mode: - queries: query - id: ID - output: - - AI Mode Results - - AI Mode Text - prompt: Find primary sources for this topic. - n_results: 4 - country: gb - language: en - device: mobile - no_cache: true - """ - - result = wrangles.recipe.run(recipe, dataframe=data) - - assert calls[0][0] == [ - "Manufacturer: WESTFALIA\nPotential part code: DN65\nDescription: union nut", - "general topic one", - "general topic two", - ] - assert calls[0][1]["prompt"] == "Find primary sources for this topic." - assert calls[0][1]["country"] == "gb" - assert calls[0][1]["device"] == "mobile" - assert result.loc[0, "AI Mode Results"][0]["search_results"][0]["input_row_id"] == "westfalia" - assert [ - payload["search_metadata"]["query_index"] - for payload in result.loc[1, "AI Mode Results"] - ] == [1, 2] - assert result.loc[2, "AI Mode Results"] == [] - assert result.loc[3, "AI Mode Results"] == [] - assert result.loc[2, "AI Mode Text"] == "" - assert "Query 1" in result.loc[0, "AI Mode Text"] - assert "A cited answer." in result.loc[0, "AI Mode Text"] - - -def test_recipe_ai_mode_supports_multiple_query_columns(mocker): - mocker.patch.object( - recipe_search._search_core, - "ai_mode", - side_effect=lambda queries, **kwargs: [ - normalized_payload(query, index) - for index, query in enumerate(queries, start=1) - ], - ) - data = pd.DataFrame( - { - "product_query": [PRODUCT_QUERY], - "general_query": ["History of ball bearings"], - "ID": [7], - } - ) - - result = recipe_search.ai_mode( - data, - queries=["product_query", "general_query"], - id="ID", - output=["product_results", "general_results"], - api_key="key", - ) - - assert result["product_results"][0][0]["search_metadata"]["query"] == PRODUCT_QUERY - assert result["general_results"][0][0]["search_metadata"]["query"] == "History of ball bearings" - - -def test_ai_mode_formatter_and_classic_formatter_contract(): - payload = normalized_payload(PRODUCT_QUERY) - payload["search_results"][0]["pricing"] = { - "price": 12.5, - "currency": "USD", - "availability": "In stock", - "vendor": "Example", - } - - text = wrangles_format.ai_mode_results_to_text([payload]) - classic_text = wrangles_format.raw_search_results_to_text([payload]) - - assert f"Query 1: {PRODUCT_QUERY}" in text - assert "Status: Success" in text - assert "A cited answer." in text - assert "Source 1" in text - assert "USD 12.5" in text - assert "## Query 1:" in classic_text - assert "Status: Success" not in classic_text - - failure = normalized_payload("failed query") - failure.update(status="Failure", error="API limit reached", search_results=[]) - failure_text = wrangles_format.ai_mode_results_to_text([failure]) - assert "Status: Failure" in failure_text - assert "Error: API limit reached" in failure_text - - -def test_ai_mode_recipe_schema_documents_supported_parameters_only(): - schema = yaml.safe_load(recipe_search.ai_mode.__doc__) - - assert schema["additionalProperties"] is False - assert schema["required"] == ["queries", "id", "output"] - assert set(schema["properties"]) == { - "queries", - "id", - "output", - "client", - "api_key", - "prompt", - "n_results", - "threads", - "country", - "language", - "location", - "uule", - "device", - "no_cache", - "include_raw_response", - } - assert schema["properties"]["device"]["enum"] == ["desktop", "tablet", "mobile"] diff --git a/tests/test_wrangles.py b/tests/test_wrangles.py index 1276e2505..977d6c127 100644 --- a/tests/test_wrangles.py +++ b/tests/test_wrangles.py @@ -930,4 +930,22 @@ def test_compare_overlap_exact_match_custom(): Test compare.overlap with exact_match parameter """ result = wrangles.compare.overlap([['test', 'test']], exact_match='MATCH') - assert result == ['MATCH'] \ No newline at end of file + assert result == ['MATCH'] + + +def test_search_ai_mode(): + query = "SKF 6205-2RS deep groove ball bearing specifications" + + result = wrangles.search.ai_mode( + query, + include_raw_response=True, + ) + + assert result["status"] == "Success", result["error"] + assert result["search_metadata"]["query"] == query + assert isinstance(result["raw_response"], dict) + assert isinstance(result["search_results"], list) + assert ( + result["extracted_content"]["answer_markdown"] + or result["extracted_content"]["text_blocks"] + ) diff --git a/wrangles/clients/serp_api.py b/wrangles/clients/serp_api.py index b060e8c01..ba66839d5 100644 --- a/wrangles/clients/serp_api.py +++ b/wrangles/clients/serp_api.py @@ -2,6 +2,7 @@ import json as _json import math as _math import re +from collections.abc import Mapping as _Mapping from typing import Union as _Union # Import our new core web helpers @@ -256,7 +257,7 @@ def _ai_mode_result(item: dict, result_type: str, query_index: int | None) -> di return result -def _normalize_ai_mode_results(response: dict, query_index: int | None, n_results: int) -> list[dict]: +def _normalize_ai_mode_results(response: dict, query_index: int | None) -> list[dict]: records = [] seen = set() sections = ( @@ -282,9 +283,9 @@ def _normalize_ai_mode_results(response: dict, query_index: int | None, n_result seen.add(key) records.append(record) - for rank, record in enumerate(records[:n_results], start=1): + for rank, record in enumerate(records, start=1): record["google_rank"] = rank - return records[:n_results] + return records class SerpApiWranglesClient: @@ -419,15 +420,13 @@ def ai_mode_single( self, query, prompt: str | None = None, - n_results: int = 10, query_index: int | None = None, country: str = "us", language: str = "en", location: str | None = None, - uule: str | None = None, - device: str = "desktop", no_cache: bool = False, include_raw_response: bool = False, + **kwargs, ) -> dict: """Perform one Google AI Mode search and normalize the provider response.""" if _is_blank_query(query): @@ -444,24 +443,24 @@ def ai_mode_single( query_text = str(query).strip() request_query = f"{prompt.strip()}\n\nQuery/product evidence:\n{query_text}" if prompt else query_text params = { + **kwargs, "engine": "google_ai_mode", "q": request_query, "output": "json", "gl": country, "hl": language, - "device": device, + "device": "desktop", "no_cache": no_cache, } if location: params["location"] = location - if uule: - params["uule"] = uule try: client = self.client_class(api_key=self.api_key) response = client.search(params) - if not isinstance(response, dict): + if not isinstance(response, _Mapping): raise TypeError("SerpAPI returned a non-object response") + response = dict(response) meta_raw = response.get("search_metadata") or {} search_params = response.get("search_parameters") or {} @@ -493,7 +492,6 @@ def ai_mode_single( search_results=[] if failed else _normalize_ai_mode_results( response, query_index, - n_results, ), answer_markdown=None if failed else response.get("reconstructed_markdown"), text_blocks=[] if failed else response.get("text_blocks"), @@ -518,15 +516,13 @@ def ai_mode_batch( self, input_data: _Union[str, list], prompt: str | None = None, - n_results: int = 10, threads: int = 10, country: str = "us", language: str = "en", location: str | None = None, - uule: str | None = None, - device: str = "desktop", no_cache: bool = False, include_raw_response: bool = False, + **kwargs, ) -> _Union[dict, list]: """Perform ordered Google AI Mode searches in parallel.""" input_was_scalar = not isinstance(input_data, list) @@ -538,15 +534,13 @@ def ai_mode_batch( lambda item: self.ai_mode_single( query=item[1], prompt=prompt, - n_results=n_results, query_index=item[0], country=country, language=language, location=location, - uule=uule, - device=device, no_cache=no_cache, include_raw_response=include_raw_response, + **kwargs, ), indexed, )) diff --git a/wrangles/recipe_wrangles/search.py b/wrangles/recipe_wrangles/search.py index 074f122d2..1502b021e 100644 --- a/wrangles/recipe_wrangles/search.py +++ b/wrangles/recipe_wrangles/search.py @@ -179,13 +179,10 @@ def ai_mode( client: str = "serpapi", api_key: str | None = None, prompt: str | None = None, - n_results: int = 10, threads: int = 10, country: str = "us", language: str = "en", location: str | None = None, - uule: str | None = None, - device: str = "desktop", no_cache: bool = False, include_raw_response: bool = False, ) -> _pd.DataFrame: @@ -223,11 +220,6 @@ def ai_mode( prompt: type: string description: Optional instruction replacing the default industrial-product research prompt. - n_results: - type: integer - minimum: 1 - description: Maximum normalized source records per query, applied after deduplication. - default: 10 threads: type: integer minimum: 1 @@ -243,18 +235,7 @@ def ai_mode( default: en location: type: string - description: Search location. Cannot be combined with uule. - uule: - type: string - description: Encoded Google location. Cannot be combined with location. - device: - type: string - description: Device type for the AI Mode request. - enum: - - desktop - - tablet - - mobile - default: desktop + description: Search location. no_cache: type: boolean description: Request a fresh SerpAPI response instead of cached results. @@ -315,13 +296,10 @@ def _to_query_list(value) -> list[str]: client=client, api_key=api_key, prompt=prompt, - n_results=n_results, threads=threads, country=country, language=language, location=location, - uule=uule, - device=device, no_cache=no_cache, include_raw_response=include_raw_response, ) diff --git a/wrangles/search.py b/wrangles/search.py index a40aa81a4..24597fea9 100644 --- a/wrangles/search.py +++ b/wrangles/search.py @@ -85,25 +85,17 @@ def ai_mode( client: str = "serpapi", api_key: str | None = None, prompt: str | None = None, - n_results: int = 10, threads: int = 10, country: str = "us", language: str = "en", location: str | None = None, - uule: str | None = None, - device: str = "desktop", no_cache: bool = False, include_raw_response: bool = False, + **kwargs, ) -> dict | list: """Search and synthesize cited content with SerpAPI Google AI Mode.""" - if not isinstance(n_results, int) or isinstance(n_results, bool) or n_results < 1: - raise ValueError("n_results must be at least 1") if not isinstance(threads, int) or isinstance(threads, bool) or threads < 1: raise ValueError("threads must be at least 1") - if location and uule: - raise ValueError("location and uule cannot both be provided") - if device not in ("desktop", "tablet", "mobile"): - raise ValueError("device must be one of: desktop, tablet, mobile") is_scalar = not isinstance(queries, list) query_list = [queries] if is_scalar else queries @@ -121,15 +113,13 @@ def ai_mode( return search_client.ai_mode_batch( queries, prompt=DEFAULT_AI_MODE_PROMPT if prompt is None else prompt, - n_results=n_results, threads=threads, country=country, language=language, location=location, - uule=uule, - device=device, no_cache=no_cache, include_raw_response=include_raw_response, + **kwargs, )