| 1 | """Entity extraction from initial search results for supplemental searches.""" |
| 2 | |
| 3 | import re |
| 4 | from collections import Counter |
| 5 | from typing import Any, Dict, List |
| 6 | |
| 7 | # Handles that appear too frequently to be useful for targeted search. |
| 8 | # These are generic/platform accounts, not topic-specific voices. |
| 9 | GENERIC_HANDLES = { |
| 10 | "elonmusk", "openai", "google", "microsoft", "apple", "meta", |
| 11 | "github", "youtube", "x", "twitter", "reddit", "wikipedia", |
| 12 | "nytimes", "washingtonpost", "cnn", "bbc", "reuters", |
| 13 | "verified", "jack", "sundarpichai", |
| 14 | } |
| 15 | |
| 16 | ENTITY_STOPWORDS = frozenset({ |
| 17 | "the", "a", "an", "to", "for", "how", "is", "in", "of", "on", "and", |
| 18 | "with", "from", "by", "at", "this", "that", "it", "what", "are", "do", |
| 19 | "can", "his", "her", "he", "she", "its", "was", "has", "new", "just", |
| 20 | "says", "said", "will", "about", "after", "now", "all", "been", "here", |
| 21 | "not", "out", "up", "more", "also", "but", "who", "year", "first", |
| 22 | "make", "being", "making", "over", "into", "than", "they", "their", |
| 23 | "would", "could", "get", "got", "some", "like", "back", "going", |
| 24 | "breaking", "https", "http", "www", "com", |
| 25 | }) |
| 26 | |
| 27 | |
| 28 | def has_anchor_signal(word: str) -> bool: |
| 29 | """True when a word carries an anchor signal: leading capital, all-caps, |
| 30 | or any digit (product/person/version anchors).""" |
| 31 | return word[0].isupper() or word.isupper() or any(char.isdigit() for char in word) |
| 32 | |
| 33 | |
| 34 | def extract_text_entities(text: str) -> set[str]: |
| 35 | """Extract significant words used by clustering and eval scoring.""" |
| 36 | words = re.sub(r"[^\w\s]", " ", text).split() |
| 37 | entities = set() |
| 38 | for word in words: |
| 39 | lower = word.lower() |
| 40 | if lower in ENTITY_STOPWORDS or len(word) <= 2: |
| 41 | continue |
| 42 | if has_anchor_signal(word) or len(word) >= 4: |
| 43 | entities.add(lower) |
| 44 | return entities |
| 45 | |
| 46 | |
| 47 | def entity_overlap(entities_a: set[str], entities_b: set[str]) -> float: |
| 48 | """Return overlap coefficient for two extracted entity sets.""" |
| 49 | if not entities_a or not entities_b: |
| 50 | return 0.0 |
| 51 | return len(entities_a & entities_b) / min(len(entities_a), len(entities_b)) |
| 52 | |
| 53 | |
| 54 | def extract_entities( |
| 55 | reddit_items: List[Dict[str, Any]], |
| 56 | x_items: List[Dict[str, Any]], |
| 57 | max_handles: int = 5, |
| 58 | max_hashtags: int = 3, |
| 59 | max_subreddits: int = 5, |
| 60 | ) -> Dict[str, List[str]]: |
| 61 | """Extract key entities from Phase 1 results for supplemental searches. |
| 62 | |
| 63 | Parses X results for @handles and #hashtags, Reddit results for subreddit |
| 64 | names and cross-referenced communities. |
| 65 | |
| 66 | Args: |
| 67 | reddit_items: Raw Reddit item dicts from Phase 1 |
| 68 | x_items: Raw X item dicts from Phase 1 |
| 69 | max_handles: Maximum handles to return |
| 70 | max_hashtags: Maximum hashtags to return |
| 71 | max_subreddits: Maximum subreddits to return |
| 72 | |
| 73 | Returns: |
| 74 | Dict with keys: x_handles, x_hashtags, reddit_subreddits |
| 75 | """ |
| 76 | handles = _extract_x_handles(x_items) |
| 77 | hashtags = _extract_x_hashtags(x_items) |
| 78 | subreddits = _extract_subreddits(reddit_items) |
| 79 | |
| 80 | return { |
| 81 | "x_handles": handles[:max_handles], |
| 82 | "x_hashtags": hashtags[:max_hashtags], |
| 83 | "reddit_subreddits": subreddits[:max_subreddits], |
| 84 | } |
| 85 | |
| 86 | |
| 87 | def _extract_x_handles(x_items: List[Dict[str, Any]]) -> List[str]: |
| 88 | """Extract and rank @handles from X results. |
| 89 | |
| 90 | Sources handles from: |
| 91 | 1. author_handle field (who posted) |
| 92 | 2. @mentions in post text (who they're talking about/to) |
| 93 | |
| 94 | Returns handles ranked by frequency, filtered for generic accounts. |
| 95 | """ |
| 96 | handle_counts = Counter() |
| 97 | |
| 98 | for item in x_items: |
| 99 | # Author handle |
| 100 | author = item.get("author_handle", "").strip().lstrip("@").lower() |
| 101 | if author and author not in GENERIC_HANDLES: |
| 102 | handle_counts[author] += 1 |
| 103 | |
| 104 | # @mentions in text |
| 105 | text = item.get("text", "") |
| 106 | mentions = re.findall(r'@(\w{1,15})', text) |
| 107 | for mention in mentions: |
| 108 | mention_lower = mention.lower() |
| 109 | if mention_lower not in GENERIC_HANDLES: |
| 110 | handle_counts[mention_lower] += 1 |
| 111 | |
| 112 | # Return all handles ranked by frequency |
| 113 | return [h for h, _ in handle_counts.most_common()] |
| 114 | |
| 115 | |
| 116 | def _extract_x_hashtags(x_items: List[Dict[str, Any]]) -> List[str]: |
| 117 | """Extract and rank #hashtags from X results. |
| 118 | |
| 119 | Returns hashtags ranked by frequency. |
| 120 | """ |
| 121 | hashtag_counts = Counter() |
| 122 | |
| 123 | for item in x_items: |
| 124 | text = item.get("text", "") |
| 125 | tags = re.findall(r'#(\w{2,30})', text) |
| 126 | for tag in tags: |
| 127 | hashtag_counts[tag.lower()] += 1 |
| 128 | |
| 129 | # Return all hashtags ranked by frequency |
| 130 | return [f"#{t}" for t, _ in hashtag_counts.most_common()] |
| 131 | |
| 132 | |
| 133 | def _extract_subreddits(reddit_items: List[Dict[str, Any]]) -> List[str]: |
| 134 | """Extract and rank subreddits from Reddit results. |
| 135 | |
| 136 | Sources from: |
| 137 | 1. subreddit field on each result |
| 138 | 2. Cross-references in comment text (e.g., "check out r/localLLaMA") |
| 139 | |
| 140 | Returns subreddits ranked by frequency. |
| 141 | """ |
| 142 | sub_counts = Counter() |
| 143 | |
| 144 | for item in reddit_items: |
| 145 | # Primary subreddit |
| 146 | sub = item.get("subreddit", "").strip().removeprefix("r/") |
| 147 | if sub: |
| 148 | sub_counts[sub] += 1 |
| 149 | |
| 150 | # Cross-references in comment insights |
| 151 | for insight in item.get("comment_insights", []): |
| 152 | cross_refs = re.findall(r'r/(\w{2,30})', insight) |
| 153 | for ref in cross_refs: |
| 154 | sub_counts[ref] += 1 |
| 155 | |
| 156 | # Cross-references in top comments |
| 157 | for comment in item.get("top_comments", []): |
| 158 | excerpt = comment.get("excerpt", "") |
| 159 | cross_refs = re.findall(r'r/(\w{2,30})', excerpt) |
| 160 | for ref in cross_refs: |
| 161 | sub_counts[ref] += 1 |
| 162 | |
| 163 | # Return subreddits ranked by frequency |
| 164 | return [sub for sub, _ in sub_counts.most_common()] |
| 165 |