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..fb81bdeb2 --- /dev/null +++ b/docs/search-ai-mode.md @@ -0,0 +1,174 @@ +# 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 + 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", + 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. | +| `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. | +| `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 may be passed as keyword arguments through the direct +Python API. AI Mode requests always use the supported desktop device. + +## 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 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. + +## 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. + +## 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, 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/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/__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..ba66839d5 100644 --- a/wrangles/clients/serp_api.py +++ b/wrangles/clients/serp_api.py @@ -1,5 +1,8 @@ import concurrent.futures as _futures +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 @@ -94,6 +97,197 @@ 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 + 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"), + ("£", "GBP"), + ("€", "EUR"), + ): + if token in value_upper: + 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) -> 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 + dedupe_link = record["link"] + if "://" not in dedupe_link: + dedupe_link = f"https://{dedupe_link}" + key = ( + _web.normalize_site(dedupe_link).lower().rstrip("/"), + record["title"].strip().lower(), + ) + if key in seen: + continue + seen.add(key) + records.append(record) + + for rank, record in enumerate(records, start=1): + record["google_rank"] = rank + return records + + class SerpApiWranglesClient: def __init__(self, api_key: str = None): if not api_key or str(api_key).strip().lower() in ("", "none", "null"): @@ -220,4 +414,135 @@ 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, + query_index: int | None = None, + country: str = "us", + language: str = "en", + location: str | None = None, + 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): + 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 = { + **kwargs, + "engine": "google_ai_mode", + "q": request_query, + "output": "json", + "gl": country, + "hl": language, + "device": "desktop", + "no_cache": no_cache, + } + if location: + params["location"] = location + + try: + client = self.client_class(api_key=self.api_key) + response = client.search(params) + 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 {} + 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_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, + } + 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, + ), + 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, + threads: int = 10, + country: str = "us", + language: str = "en", + location: str | None = None, + 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) + 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, + query_index=item[0], + country=country, + language=language, + location=location, + no_cache=no_cache, + include_raw_response=include_raw_response, + **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..1502b021e 100644 --- a/wrangles/recipe_wrangles/search.py +++ b/wrangles/recipe_wrangles/search.py @@ -171,6 +171,188 @@ 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, + threads: int = 10, + country: str = "us", + language: str = "en", + location: str | None = None, + 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. + 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. + 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, + threads=threads, + country=country, + language=language, + location=location, + 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..24597fea9 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,49 @@ def find_links( ) +def ai_mode( + queries: str | list, + client: str = "serpapi", + api_key: str | None = None, + prompt: str | None = None, + threads: int = 10, + country: str = "us", + language: str = "en", + location: str | None = None, + 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(threads, int) or isinstance(threads, bool) or threads < 1: + raise ValueError("threads must be at least 1") + + 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, + threads=threads, + country=country, + language=language, + location=location, + no_cache=no_cache, + include_raw_response=include_raw_response, + **kwargs, + ) + + def retrieve_link_content( urls: str | list, client: str = "google_url_context",