segment_classifier 0.1.0__py3-none-any.whl

segment_classifier/stages/fuzzy_cluster.py

@@ -0,0 +1,101 @@
+from sklearn.feature_extraction.text import TfidfVectorizer
+from segment_classifier.cache.l2_cache import L2FuzzyCache
+from segment_classifier.config import CacheConfig
+from segment_classifier.models import (
+    InputSegment, ClassifiedSegment, ClassificationStage, ComponentType
+)
+from segment_classifier.utils.html_normalizer import NormalizedSegment
+
+
+class FuzzyClusterStage:
+    def __init__(self, cache: L2FuzzyCache, config: CacheConfig):
+        self.cache = cache
+        self.config = config
+
+        # Use TfidfVectorizer, but fit it once on a known dummy vocabulary
+        # or just use sklearn's HashingVectorizer which is stateless.
+        # But since prompt specifies "use sklearn TfidfVectorizer", we must pre-fit it
+        # or use HashingVectorizer with TfidfTransformer. HashingVectorizer directly
+        # isn't TfidfVectorizer. Let's stick to HashingVectorizer + tfidf for stability?
+        # Actually, prompt says "use sklearn TfidfVectorizer". We can use TfidfVectorizer
+        # but with a fixed vocabulary if we hash. Or we can just build a stateless tfidf
+        # using HashingVectorizer + TfidfTransformer.
+        # But wait, code review said: "You must either use a HashingVectorizer... or pre-fit the vectorizer".
+        # Let's use HashingVectorizer + TfidfTransformer to get TF-IDF scaling on fixed dimensions.
+        from sklearn.feature_extraction.text import HashingVectorizer, TfidfTransformer
+        from sklearn.pipeline import Pipeline
+        self.vectorizer = Pipeline([
+            ('hash', HashingVectorizer(
+                analyzer='char',
+                ngram_range=(2, 4),
+                n_features=512,
+                lowercase=True,
+                norm=None # Let TfidfTransformer handle normalization
+            )),
+            ('tfidf', TfidfTransformer())
+        ])
+        # Fit once on empty/dummy string to initialize tfidf idf_ to smooth values
+        self.vectorizer.fit(["dummy"])
+        self._is_fitted = True
+
+    def _build_fingerprint_string(self, normalized: NormalizedSegment) -> str:
+        return f"{normalized.skeleton} {normalized.attrs_fingerprint} {' '.join(normalized.class_tokens)}"
+
+    def _vectorize(self, fingerprint_string: str) -> list[float]:
+        matrix = self.vectorizer.transform([fingerprint_string])
+        return matrix.toarray()[0].tolist()
+
+    async def classify(
+        self,
+        segment: InputSegment,
+        normalized: NormalizedSegment,
+        fingerprint_hash: str,
+    ) -> ClassifiedSegment | None:
+        """
+        1. Build fingerprint string
+        2. Vectorize
+        3. find_nearest in L2Cache
+        4. Return ClassifiedSegment (stage=L2_FUZZY_CACHE) or None
+        """
+        fingerprint_string = self._build_fingerprint_string(normalized)
+        vector = self._vectorize(fingerprint_string)
+
+        nearest = await self.cache.find_nearest(vector, self.config.l2_similarity_threshold)
+        if nearest:
+            # Penalty of 0.05 for fuzzy
+            confidence = max(0.0, nearest.confidence - 0.05)
+
+            return ClassifiedSegment(
+                segment_id=segment.segment_id,
+                page_url=segment.page_url,
+                page_slug=segment.page_slug,
+                raw_html=segment.raw_html,
+                text_content=segment.text_content,
+                position_hint=segment.position_hint,
+                component_type=nearest.component_type,
+                classification_stage=ClassificationStage.L2_FUZZY_CACHE,
+                confidence=confidence,
+                fingerprint_hash=fingerprint_hash,
+                cluster_id=nearest.cluster_id
+            )
+
+        return None
+
+    async def register(
+        self,
+        fingerprint_hash: str,
+        normalized: NormalizedSegment,
+        component_type: ComponentType,
+        confidence: float,
+    ) -> None:
+        """
+        Called after LLM resolves a segment to register it in L2 for future lookup.
+        """
+        fingerprint_string = self._build_fingerprint_string(normalized)
+        vector = self._vectorize(fingerprint_string)
+
+        nearest = await self.cache.find_nearest(vector, self.config.l2_similarity_threshold)
+        if nearest:
+            await self.cache.add_to_cluster(nearest.cluster_id, fingerprint_hash, vector)
+        else:
+            await self.cache.create_cluster(fingerprint_hash, vector, component_type, confidence)

segment_classifier/stages/llm_classifier.py

@@ -0,0 +1,271 @@
+import asyncio
+import json
+import logging
+import litellm
+from typing import Any
+from segment_classifier.models import (
+    InputSegment, ClassifiedSegment, LLMClassificationRequest,
+    LLMClassificationResult, ClassificationStage, ComponentType
+)
+from segment_classifier.utils.html_normalizer import NormalizedSegment
+from segment_classifier.config import ClassifierSettings, ModelFeatureConfig
+
+logger = logging.getLogger(__name__)
+
+
+class LLMBatchClassifier:
+    def __init__(self, settings: ClassifierSettings):
+        self.settings = settings
+        self._semaphore = asyncio.Semaphore(
+            settings.litellm_max_concurrent_batches
+        )
+        self._model_usage: dict[str, int] = {}
+
+        # Set api key globally or per call, litellm supports both
+        if settings.litellm_api_key:
+            litellm.api_key = settings.litellm_api_key
+
+    def select_model(
+        self,
+        normalized: NormalizedSegment,
+        segment: InputSegment,
+    ) -> str:
+        """Feature-based model routing."""
+        cfg = self.settings.model_routing
+
+        # 1. High complexity
+        if (normalized.dom_depth > cfg.high_complexity_dom_depth_threshold or
+            normalized.unique_tag_count > cfg.high_complexity_unique_tag_threshold or
+            segment.sibling_count == 0):
+            return cfg.high_complexity_model
+
+        # 2. Fast
+        elif (normalized.dom_depth <= cfg.fast_model_max_dom_depth and
+              normalized.unique_tag_count <= 3):
+            return cfg.fast_model
+
+        # 3. Standard
+        else:
+            return cfg.standard_model
+
+    def _build_request(
+        self,
+        segment: InputSegment,
+        normalized: NormalizedSegment,
+        fingerprint_hash: str,
+    ) -> LLMClassificationRequest:
+        """Construct LLMClassificationRequest from segment + normalized data."""
+        return LLMClassificationRequest(
+            segment_id=segment.segment_id,
+            fingerprint_hash=fingerprint_hash,
+            normalized_html=normalized.skeleton,
+            position_hint=segment.position_hint,
+            sibling_count=segment.sibling_count,
+            url_hints=segment.url_path_segments,
+            dom_depth=normalized.dom_depth,
+            child_tag_counts=normalized.child_tag_counts,
+            text_density_ratio=normalized.text_density_ratio
+        )
+
+    async def _call_litellm(
+        self,
+        model: str,
+        requests: list[LLMClassificationRequest],
+    ) -> list[LLMClassificationResult]:
+        """
+        Make one LiteLLM acompletion call for a batch.
+        Parse response. Return list of results.
+        On error: return UNKNOWN for all items.
+        """
+        prompt = """You are an expert UI component classifier. Given a list of HTML segment descriptors, classify each into exactly one component type.
+
+Available component types:
+"""
+        prompt += ", ".join([c.value for c in ComponentType]) + "\n\n"
+        prompt += """For each segment respond with valid JSON array only (no markdown, no explanation):
+[
+  {
+    "segment_id": "...",
+    "component_type": "...",
+    "confidence": 0.0-1.0,
+    "reasoning": "one sentence"
+  }
+]
+
+Rules:
+- Use normalized_html structure only, ignore content values
+- sibling_count >= 3 strongly suggests a collection item
+- position_hint=top/bottom suggests layout components
+- url_hints provide page context
+- Respond ONLY with the JSON array. No preamble."""
+
+        user_content = f"Classify these {len(requests)} segments:\n"
+        user_content += json.dumps([r.model_dump() for r in requests], indent=2)
+
+        messages = [
+            {"role": "system", "content": prompt},
+            {"role": "user", "content": user_content}
+        ]
+
+        try:
+            response = await litellm.acompletion(
+                model=model,
+                messages=messages,
+                timeout=self.settings.litellm_timeout_seconds,
+            )
+
+            # Record usage
+            self._model_usage[model] = self._model_usage.get(model, 0) + 1
+
+            raw_response = response.choices[0].message.content
+            # Strip markdown
+            raw_response = raw_response.strip()
+            if raw_response.startswith("```json"):
+                raw_response = raw_response[7:]
+            elif raw_response.startswith("```"):
+                raw_response = raw_response[3:]
+            if raw_response.endswith("```"):
+                raw_response = raw_response[:-3]
+
+            parsed = json.loads(raw_response.strip())
+
+            results = []
+            for item in parsed:
+                try:
+                    results.append(LLMClassificationResult.model_validate(item))
+                except Exception as e:
+                    logger.warning(f"Error validating LLM response item: {e}")
+                    results.append(LLMClassificationResult(
+                        segment_id=item.get("segment_id", ""),
+                        component_type=ComponentType.UNKNOWN,
+                        confidence=0.0,
+                        reasoning=f"Parse error: {e}"
+                    ))
+
+            # Ensure all segments are accounted for
+            parsed_ids = {r.segment_id for r in results}
+            for req in requests:
+                if req.segment_id not in parsed_ids:
+                    results.append(LLMClassificationResult(
+                        segment_id=req.segment_id,
+                        component_type=ComponentType.UNKNOWN,
+                        confidence=0.0,
+                        reasoning="Missing from LLM response"
+                    ))
+
+            return results
+
+        except Exception as e:
+            logger.warning(f"LLM call failed: {e}")
+            self._model_usage[model] = self._model_usage.get(model, 0) + 1 # Still count as call
+            return [
+                LLMClassificationResult(
+                    segment_id=r.segment_id,
+                    component_type=ComponentType.UNKNOWN,
+                    confidence=0.0,
+                    reasoning=f"LLM Error: {e}"
+                )
+                for r in requests
+            ]
+
+    async def classify_batch(
+        self,
+        items: list[tuple[InputSegment, NormalizedSegment, str]],
+    ) -> list[ClassifiedSegment]:
+        """
+        1. Group items by selected model
+        2. Split each group into sub-batches of litellm_batch_size
+        3. asyncio.gather all sub-batches under semaphore
+        4. Return flat list of ClassifiedSegments
+        """
+        # Group by model
+        grouped: dict[str, list[tuple[InputSegment, NormalizedSegment, str]]] = {}
+        for item in items:
+            segment, normalized, _ = item
+            model = self.select_model(normalized, segment)
+            if model not in grouped:
+                grouped[model] = []
+            grouped[model].append(item)
+
+        all_classified: list[ClassifiedSegment] = []
+
+        async def process_subbatch(model_name: str, subbatch: list[tuple[InputSegment, NormalizedSegment, str]]):
+            async with self._semaphore:
+                requests = [self._build_request(s, n, h) for s, n, h in subbatch]
+                results = await self._call_litellm(model_name, requests)
+
+                # Map results back
+                result_map = {r.segment_id: r for r in results}
+
+                for segment, normalized, fp_hash in subbatch:
+                    res = result_map.get(segment.segment_id)
+                    if res:
+                        all_classified.append(
+                            ClassifiedSegment(
+                                segment_id=segment.segment_id,
+                                page_url=segment.page_url,
+                                page_slug=segment.page_slug,
+                                raw_html=segment.raw_html,
+                                text_content=segment.text_content,
+                                position_hint=segment.position_hint,
+                                component_type=res.component_type,
+                                classification_stage=ClassificationStage.LLM,
+                                confidence=res.confidence,
+                                fingerprint_hash=fp_hash,
+                                llm_model_used=model_name,
+                                llm_raw_response=res.reasoning # Hack: put reasoning here
+                            )
+                        )
+                    else:
+                        all_classified.append(
+                            ClassifiedSegment(
+                                segment_id=segment.segment_id,
+                                page_url=segment.page_url,
+                                page_slug=segment.page_slug,
+                                raw_html=segment.raw_html,
+                                text_content=segment.text_content,
+                                position_hint=segment.position_hint,
+                                component_type=ComponentType.UNKNOWN,
+                                classification_stage=ClassificationStage.LLM,
+                                confidence=0.0,
+                                fingerprint_hash=fp_hash,
+                                llm_model_used=model_name
+                            )
+                        )
+
+        tasks = []
+        for model_name, group_items in grouped.items():
+            # Split into sub-batches
+            for i in range(0, len(group_items), self.settings.litellm_batch_size):
+                subbatch = group_items[i:i + self.settings.litellm_batch_size]
+                tasks.append(process_subbatch(model_name, subbatch))
+
+        await asyncio.gather(*tasks)
+
+        # Ensure the returned list is in the exact order as the input items
+        result_map = {res.segment_id: res for res in all_classified}
+        ordered_classified = []
+        for segment, _, _ in items:
+            if segment.segment_id in result_map:
+                ordered_classified.append(result_map[segment.segment_id])
+            else:
+                # Fallback, though process_subbatch should populate it
+                ordered_classified.append(
+                    ClassifiedSegment(
+                        segment_id=segment.segment_id,
+                        page_url=segment.page_url,
+                        page_slug=segment.page_slug,
+                        raw_html=segment.raw_html,
+                        text_content=segment.text_content,
+                        position_hint=segment.position_hint,
+                        component_type=ComponentType.UNKNOWN,
+                        classification_stage=ClassificationStage.LLM,
+                        confidence=0.0,
+                        fingerprint_hash=""
+                    )
+                )
+        return ordered_classified
+
+    @property
+    def model_usage(self) -> dict[str, int]:
+        return dict(self._model_usage)

segment_classifier/stages/rule_based.py

@@ -0,0 +1,287 @@
+from dataclasses import dataclass
+from typing import Callable
+from segment_classifier.models import (
+    InputSegment, ComponentType, ClassificationStage, ClassifiedSegment, SegmentPosition
+)
+from segment_classifier.utils.html_normalizer import NormalizedSegment
+
+
+@dataclass
+class ClassificationRule:
+    name: str
+    condition: Callable[[InputSegment, NormalizedSegment], bool]
+    component_type: ComponentType
+    confidence: float
+    priority: int = 50
+
+
+class RuleBasedClassifier:
+    """
+    Apply ordered classification rules to a segment.
+    Rules evaluated highest priority first.
+    Returns ClassifiedSegment with stage=RULE_BASED or None if no match.
+    """
+
+    def __init__(self, confidence_threshold: float = 0.90):
+        self.confidence_threshold = confidence_threshold
+        self.rules: list[ClassificationRule] = self._build_rules()
+
+    def _build_rules(self) -> list[ClassificationRule]:
+        """
+        Instantiate all rules described above.
+        Sort by priority descending before returning.
+        """
+        rules = [
+            # Position rules
+            ClassificationRule(
+                name="top_header_nav",
+                condition=lambda s, n: s.position_hint == SegmentPosition.TOP and n.root_tag in {"header", "nav"},
+                component_type=ComponentType.LAYOUT_HEADER,
+                confidence=0.97,
+                priority=100
+            ),
+            ClassificationRule(
+                name="bottom_footer",
+                condition=lambda s, n: s.position_hint == SegmentPosition.BOTTOM and n.root_tag in {"footer"},
+                component_type=ComponentType.LAYOUT_FOOTER,
+                confidence=0.97,
+                priority=99
+            ),
+            ClassificationRule(
+                name="top_has_nav",
+                condition=lambda s, n: s.position_hint == SegmentPosition.TOP and n.child_tag_counts.get("nav", 0) > 0,
+                component_type=ComponentType.LAYOUT_NAV,
+                confidence=0.93,
+                priority=98
+            ),
+
+            # Tag-based rules
+            ClassificationRule(
+                name="root_footer",
+                condition=lambda s, n: n.root_tag == "footer",
+                component_type=ComponentType.LAYOUT_FOOTER,
+                confidence=0.99,
+                priority=90
+            ),
+            ClassificationRule(
+                name="root_header",
+                condition=lambda s, n: n.root_tag == "header",
+                component_type=ComponentType.LAYOUT_HEADER,
+                confidence=0.99,
+                priority=90
+            ),
+            ClassificationRule(
+                name="nav_element",
+                condition=lambda s, n: n.root_tag == "nav" or (n.child_tag_counts.get("nav", 0) > 0 and n.child_tag_counts.get("a", 0) > 3),
+                component_type=ComponentType.LAYOUT_NAV,
+                confidence=0.95,
+                priority=85
+            ),
+            ClassificationRule(
+                name="form_element",
+                condition=lambda s, n: n.root_tag == "form" or n.child_tag_counts.get("input", 0) > 1,
+                component_type=ComponentType.UI_FORM,
+                confidence=0.93,
+                priority=84
+            ),
+            ClassificationRule(
+                name="table_element",
+                condition=lambda s, n: n.child_tag_counts.get("table", 0) > 0,
+                component_type=ComponentType.UI_TABLE,
+                confidence=0.90,
+                priority=83
+            ),
+            ClassificationRule(
+                name="modal_element",
+                condition=lambda s, n: n.child_tag_counts.get("dialog", 0) > 0 or "modal" in n.class_tokens,
+                component_type=ComponentType.UI_MODAL,
+                confidence=0.91,
+                priority=82
+            ),
+
+            # Sibling repetition rules
+            ClassificationRule(
+                name="product_card_repetition",
+                condition=lambda s, n: s.sibling_count >= 3 and "card" in n.class_tokens,
+                component_type=ComponentType.COLLECTION_PRODUCT_CARD,
+                confidence=0.88,
+                priority=70
+            ),
+            ClassificationRule(
+                name="product_card_img_header",
+                condition=lambda s, n: s.sibling_count >= 3 and n.child_tag_counts.get("img", 0) > 0 and (n.child_tag_counts.get("h2", 0) > 0 or n.child_tag_counts.get("h3", 0) > 0),
+                component_type=ComponentType.COLLECTION_PRODUCT_CARD,
+                confidence=0.85,
+                priority=69
+            ),
+            ClassificationRule(
+                name="blog_card_repetition",
+                condition=lambda s, n: s.sibling_count >= 3 and "article" in n.class_tokens,
+                component_type=ComponentType.COLLECTION_BLOG_CARD,
+                confidence=0.87,
+                priority=68
+            ),
+            ClassificationRule(
+                name="nav_list_repetition",
+                condition=lambda s, n: s.sibling_count >= 5 and n.child_tag_counts.get("li", 0) > 4,
+                component_type=ComponentType.LAYOUT_NAV,
+                confidence=0.86,
+                priority=67
+            ),
+
+            # URL hint rules
+            ClassificationRule(
+                name="url_product_card",
+                condition=lambda s, n: any(x in s.url_path_segments for x in ["product", "shop", "store"]) and s.sibling_count >= 2,
+                component_type=ComponentType.COLLECTION_PRODUCT_CARD,
+                confidence=0.84,
+                priority=60
+            ),
+            ClassificationRule(
+                name="url_blog_card",
+                condition=lambda s, n: any(x in s.url_path_segments for x in ["blog", "post", "article"]),
+                component_type=ComponentType.COLLECTION_BLOG_CARD,
+                confidence=0.83,
+                priority=59
+            ),
+            ClassificationRule(
+                name="url_news_item",
+                condition=lambda s, n: "news" in s.url_path_segments,
+                component_type=ComponentType.COLLECTION_NEWS_ITEM,
+                confidence=0.83,
+                priority=58
+            ),
+
+            # Class token rules
+            ClassificationRule(
+                name="class_product_price",
+                condition=lambda s, n: "price" in n.class_tokens and n.child_tag_counts.get("img", 0) > 0,
+                component_type=ComponentType.COLLECTION_PRODUCT_CARD,
+                confidence=0.89,
+                priority=50
+            ),
+            ClassificationRule(
+                name="class_hero",
+                condition=lambda s, n: "hero" in n.class_tokens and n.dom_depth <= 4,
+                component_type=ComponentType.SECTION_HERO,
+                confidence=0.91,
+                priority=49
+            ),
+            ClassificationRule(
+                name="class_testimonial",
+                condition=lambda s, n: "testimonial" in n.class_tokens,
+                component_type=ComponentType.SECTION_TESTIMONIAL,
+                confidence=0.90,
+                priority=48
+            ),
+            ClassificationRule(
+                name="class_faq",
+                condition=lambda s, n: "faq" in n.class_tokens,
+                component_type=ComponentType.SECTION_FAQ,
+                confidence=0.90,
+                priority=47
+            ),
+            ClassificationRule(
+                name="class_pricing",
+                condition=lambda s, n: "pricing" in n.class_tokens,
+                component_type=ComponentType.SECTION_PRICING,
+                confidence=0.91,
+                priority=46
+            ),
+            ClassificationRule(
+                name="class_breadcrumb",
+                condition=lambda s, n: "breadcrumb" in n.class_tokens or "breadcrumb" in n.attrs_fingerprint,
+                component_type=ComponentType.LAYOUT_BREADCRUMB,
+                confidence=0.95,
+                priority=45
+            ),
+            ClassificationRule(
+                name="class_carousel",
+                condition=lambda s, n: "carousel" in n.class_tokens or n.child_tag_counts.get("swiper-slide", 0) > 0,
+                component_type=ComponentType.UI_CAROUSEL,
+                confidence=0.89,
+                priority=44
+            ),
+            ClassificationRule(
+                name="class_pagination",
+                condition=lambda s, n: "pagination" in n.class_tokens,
+                component_type=ComponentType.UI_PAGINATION,
+                confidence=0.94,
+                priority=43
+            ),
+            ClassificationRule(
+                name="class_search",
+                condition=lambda s, n: "search" in n.class_tokens and n.child_tag_counts.get("input", 0) > 0,
+                component_type=ComponentType.UI_SEARCH,
+                confidence=0.93,
+                priority=42
+            ),
+            ClassificationRule(
+                name="class_cta",
+                condition=lambda s, n: "cta" in n.class_tokens,
+                component_type=ComponentType.SECTION_CTA,
+                confidence=0.88,
+                priority=41
+            ),
+
+            # Text density rules
+            ClassificationRule(
+                name="text_article",
+                condition=lambda s, n: n.text_density_ratio > 0.6 and n.root_tag in {"article", "main"},
+                component_type=ComponentType.CONTENT_ARTICLE,
+                confidence=0.85,
+                priority=30
+            ),
+            ClassificationRule(
+                name="text_rich_text",
+                condition=lambda s, n: n.text_density_ratio > 0.5 and n.dom_depth <= 3,
+                component_type=ComponentType.CONTENT_RICH_TEXT,
+                confidence=0.80,
+                priority=29
+            ),
+
+            # Media rules
+            ClassificationRule(
+                name="media_img",
+                condition=lambda s, n: n.child_tag_counts.get("img", 0) > 3 and n.text_density_ratio < 0.1,
+                component_type=ComponentType.CONTENT_MEDIA,
+                confidence=0.87,
+                priority=20
+            ),
+            ClassificationRule(
+                name="media_video",
+                condition=lambda s, n: n.child_tag_counts.get("video", 0) > 0,
+                component_type=ComponentType.CONTENT_MEDIA,
+                confidence=0.92,
+                priority=19
+            ),
+        ]
+
+        rules.sort(key=lambda r: r.priority, reverse=True)
+        return rules
+
+    def classify(
+        self,
+        segment: InputSegment,
+        normalized: NormalizedSegment,
+    ) -> ClassifiedSegment | None:
+        """
+        Try each rule in order. Return first match above confidence_threshold.
+        Return None if no rule fires.
+        """
+        for rule in self.rules:
+            if rule.condition(segment, normalized):
+                if rule.confidence >= self.confidence_threshold:
+                    return ClassifiedSegment(
+                        segment_id=segment.segment_id,
+                        page_url=segment.page_url,
+                        page_slug=segment.page_slug,
+                        raw_html=segment.raw_html,
+                        text_content=segment.text_content,
+                        position_hint=segment.position_hint,
+                        component_type=rule.component_type,
+                        classification_stage=ClassificationStage.RULE_BASED,
+                        confidence=rule.confidence,
+                        fingerprint_hash=normalized.fingerprint_hash()
+                    )
+        return None

segment_classifier/utils/__init__.py

@@ -0,0 +1,3 @@
+from .html_normalizer import normalize_segment, NormalizedSegment
+
+__all__ = ["normalize_segment", "NormalizedSegment"]