Upload web_crawler.py with huggingface_hub

web_crawler.py

@@ -0,0 +1,140 @@
+"""
+Tavily-powered web crawler for retrieving polymer datasheets.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from tavily import TavilyClient
+
+import config
+
+logger = logging.getLogger(__name__)
+
+
+def build_search_queries(
+    manufacturer: str,
+    polymer_family: str,
+    grade: str = "",
+) -> list[str]:
+    """
+    Generate multiple targeted search queries to maximize the chance
+    of finding the correct technical datasheet.
+    """
+    base_terms = []
+    if manufacturer:
+        base_terms.append(manufacturer)
+    if polymer_family:
+        base_terms.append(polymer_family)
+    if grade:
+        base_terms.append(grade)
+
+    base = " ".join(base_terms)
+
+    queries = [
+        f"{base} technical data sheet properties",
+        f"{base} TDS material properties datasheet",
+        f"{base} mechanical thermal physical properties",
+    ]
+
+    # Add manufacturer-specific datasheet portal query
+    if manufacturer:
+        queries.append(
+            f"site:{_guess_domain(manufacturer)} {polymer_family} {grade} datasheet"
+        )
+
+    # Add aggregator queries
+    queries.append(f"{base} datasheet matweb OR omnexus OR UL Prospector")
+
+    return queries
+
+
+def _guess_domain(manufacturer: str) -> str:
+    """Attempt to guess manufacturer domain for site-scoped search."""
+    name = manufacturer.lower().replace(" ", "")
+    for domain in config.TRUSTED_DOMAINS:
+        if name in domain:
+            return domain
+    return f"{name}.com"
+
+
+def search_datasheets(
+    manufacturer: str,
+    polymer_family: str,
+    grade: str = "",
+) -> tuple[list[dict[str, Any]], str]:
+    """
+    Execute Tavily searches and return (results_list, aggregated_raw_content).
+    """
+    client = TavilyClient(api_key=config.TAVILY_API_KEY)
+
+    queries = build_search_queries(manufacturer, polymer_family, grade)
+    all_results: list[dict[str, Any]] = []
+    seen_urls: set[str] = set()
+    raw_texts: list[str] = []
+
+    for query in queries:
+        try:
+            logger.info("Searching: %s", query)
+            response = client.search(
+                query=query,
+                search_depth=config.TAVILY_SEARCH_DEPTH,
+                max_results=config.TAVILY_MAX_RESULTS,
+                include_raw_content=config.TAVILY_INCLUDE_RAW_CONTENT,
+                include_domains=config.TRUSTED_DOMAINS,
+            )
+
+            for result in response.get("results", []):
+                url = result.get("url", "")
+                if url in seen_urls:
+                    continue
+                seen_urls.add(url)
+
+                all_results.append(result)
+
+                # Collect raw content for LLM parsing
+                raw = result.get("raw_content") or result.get("content", "")
+                if raw:
+                    raw_texts.append(
+                        f"--- Source: {url} ---\n{raw[:8000]}\n"
+                    )
+
+        except Exception as exc:
+            logger.warning("Search failed for query '%s': %s", query, exc)
+
+    aggregated = "\n".join(raw_texts)
+
+    # Truncate to ~30k chars to stay within LLM context window
+    if len(aggregated) > 30_000:
+        aggregated = aggregated[:30_000] + "\n\n[Content truncated]"
+
+    logger.info(
+        "Collected %d unique results, %d chars of raw content",
+        len(all_results),
+        len(aggregated),
+    )
+
+    return all_results, aggregated
+
+
+def extract_from_url(url: str) -> tuple[list[dict[str, Any]], str]:
+    """
+    Use Tavily extract to get content from a specific URL.
+    Useful when the user provides a direct datasheet link.
+    """
+    client = TavilyClient(api_key=config.TAVILY_API_KEY)
+
+    try:
+        response = client.extract(urls=[url])
+        results = response.get("results", [])
+        raw_texts = []
+        for r in results:
+            raw = r.get("raw_content", "")
+            if raw:
+                raw_texts.append(raw[:15000])
+        return results, "\n".join(raw_texts)
+    except Exception as exc:
+        logger.error("URL extraction failed for %s: %s", url, exc)
+        return [], ""