segment_classifier 0.1.0__py3-none-any.whl

segment_classifier/__init__.py

@@ -0,0 +1,4 @@
+from .pipeline import ClassifierPipeline
+from .config import ClassifierSettings
+
+__all__ = ["ClassifierPipeline", "ClassifierSettings"]

segment_classifier/cache/__init__.py

@@ -0,0 +1,4 @@
+from .l1_cache import L1FingerprintCache
+from .l2_cache import L2FuzzyCache
+
+__all__ = ["L1FingerprintCache", "L2FuzzyCache"]

segment_classifier/cache/l1_cache.py

@@ -0,0 +1,71 @@
+import asyncio
+import json
+from pathlib import Path
+import aiofiles
+from pydantic import ValidationError
+from segment_classifier.models import FingerprintRecord
+
+
+class L1FingerprintCache:
+    def __init__(self, cache_path: str, auto_persist_every: int = 50):
+        self._path = Path(cache_path)
+        self._store: dict[str, FingerprintRecord] = {}
+        self._lock = asyncio.Lock()
+        self._write_count = 0
+        self._auto_persist_every = auto_persist_every
+
+    async def load(self) -> None:
+        if not self._path.exists():
+            return
+
+        async with aiofiles.open(self._path, "r", encoding="utf-8") as f:
+            content = await f.read()
+            if not content.strip():
+                return
+
+            try:
+                data = json.loads(content)
+                for key, val in data.items():
+                    self._store[key] = FingerprintRecord.model_validate(val)
+            except (json.JSONDecodeError, ValidationError) as e:
+                # Log error in real app, we'll just ignore and start fresh here
+                pass
+
+    async def get(self, fingerprint_hash: str) -> FingerprintRecord | None:
+        async with self._lock:
+            return self._store.get(fingerprint_hash)
+
+    async def set(self, fingerprint_hash: str, record: FingerprintRecord) -> None:
+        async with self._lock:
+            self._store[fingerprint_hash] = record
+            self._write_count += 1
+            if self._write_count >= self._auto_persist_every:
+                self._write_count = 0
+                await self._persist_unsafe()
+
+    async def increment_hit(self, fingerprint_hash: str) -> None:
+        async with self._lock:
+            record = self._store.get(fingerprint_hash)
+            if record:
+                record.hit_count += 1
+                self._store[fingerprint_hash] = record
+                self._write_count += 1
+                if self._write_count >= self._auto_persist_every:
+                    self._write_count = 0
+                    await self._persist_unsafe()
+
+    async def _persist_unsafe(self) -> None:
+        """Called internally when lock is already acquired."""
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        data = {k: v.model_dump(mode="json") for k, v in self._store.items()}
+        async with aiofiles.open(self._path, "w", encoding="utf-8") as f:
+            await f.write(json.dumps(data, indent=2))
+
+    async def persist(self) -> None:
+        async with self._lock:
+            self._write_count = 0
+            await self._persist_unsafe()
+
+    @property
+    def size(self) -> int:
+        return len(self._store)

segment_classifier/cache/l2_cache.py

@@ -0,0 +1,157 @@
+import asyncio
+import json
+import uuid
+from pathlib import Path
+import aiofiles
+import numpy as np
+from pydantic import ValidationError
+from segment_classifier.models import ClusterRecord, ComponentType
+
+
+class L2FuzzyCache:
+    def __init__(
+        self,
+        cache_path: str,
+        embeddings_path: str,
+        similarity_threshold: float = 0.85,
+        max_cluster_size: int = 50,
+        persist_on_update: bool = True
+    ):
+        self._path = Path(cache_path)
+        self._embeddings_path = Path(embeddings_path)
+        self._similarity_threshold = similarity_threshold
+        self._max_cluster_size = max_cluster_size
+        self._persist_on_update = persist_on_update
+
+        self._store: list[ClusterRecord] = []
+        # Matrix storing centroids, parallel to self._store
+        self._centroids: np.ndarray | None = None
+
+        self._lock = asyncio.Lock()
+
+    async def load(self) -> None:
+        if not self._path.exists() or not self._embeddings_path.exists():
+            return
+
+        async with aiofiles.open(self._path, "r", encoding="utf-8") as f:
+            content = await f.read()
+            if content.strip():
+                try:
+                    data = json.loads(content)
+                    self._store = [ClusterRecord.model_validate(val) for val in data]
+                except (json.JSONDecodeError, ValidationError):
+                    self._store = []
+
+        if self._store:
+            try:
+                self._centroids = np.load(str(self._embeddings_path))
+                if len(self._store) != self._centroids.shape[0]:
+                    # Mismatch between JSON and npy, reset
+                    self._store = []
+                    self._centroids = None
+            except Exception:
+                self._store = []
+                self._centroids = None
+
+    async def find_nearest(self, vector: list[float], threshold: float | None = None) -> ClusterRecord | None:
+        async with self._lock:
+            if not self._store or self._centroids is None:
+                return None
+
+            query = np.array(vector)
+            query_norm = np.linalg.norm(query)
+            if query_norm == 0:
+                return None
+
+            # Cosine similarity
+            dot_products = np.dot(self._centroids, query)
+            centroid_norms = np.linalg.norm(self._centroids, axis=1)
+            # Avoid division by zero
+            centroid_norms[centroid_norms == 0] = 1
+
+            similarities = dot_products / (centroid_norms * query_norm)
+
+            best_idx = np.argmax(similarities)
+            best_sim = similarities[best_idx]
+
+            check_threshold = threshold if threshold is not None else self._similarity_threshold
+
+            if best_sim >= check_threshold:
+                return self._store[best_idx]
+            return None
+
+    async def add_to_cluster(self, cluster_id: str, fingerprint_hash: str, vector: list[float]) -> None:
+        async with self._lock:
+            idx = next((i for i, c in enumerate(self._store) if c.cluster_id == cluster_id), None)
+            if idx is not None and self._centroids is not None:
+                cluster = self._store[idx]
+
+                # Check size
+                if len(cluster.member_fingerprints) < self._max_cluster_size:
+                    if fingerprint_hash not in cluster.member_fingerprints:
+                        cluster.member_fingerprints.append(fingerprint_hash)
+
+                        # Update centroid (running mean)
+                        n = len(cluster.member_fingerprints)
+                        old_centroid = self._centroids[idx]
+                        new_vec = np.array(vector)
+                        # (old_centroid * (n-1) + new_vec) / n
+                        new_centroid = (old_centroid * (n - 1) + new_vec) / n
+
+                        self._centroids[idx] = new_centroid
+                        cluster.centroid_vector = new_centroid.tolist()
+
+                        if self._persist_on_update:
+                            await self._persist_unsafe()
+
+    async def create_cluster(
+        self,
+        fingerprint_hash: str,
+        vector: list[float],
+        component_type: ComponentType,
+        confidence: float,
+    ) -> ClusterRecord:
+        async with self._lock:
+            cluster_id = str(uuid.uuid4())
+            record = ClusterRecord(
+                cluster_id=cluster_id,
+                centroid_vector=vector,
+                component_type=component_type,
+                confidence=confidence,
+                member_fingerprints=[fingerprint_hash]
+            )
+
+            self._store.append(record)
+
+            new_vec = np.array([vector])
+            if self._centroids is None:
+                self._centroids = new_vec
+            else:
+                self._centroids = np.vstack([self._centroids, new_vec])
+
+            if self._persist_on_update:
+                await self._persist_unsafe()
+
+            return record
+
+    async def _persist_unsafe(self) -> None:
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Write JSON
+        data = [c.model_dump(mode="json") for c in self._store]
+        async with aiofiles.open(self._path, "w", encoding="utf-8") as f:
+            await f.write(json.dumps(data, indent=2))
+
+        # Write npy
+        if self._centroids is not None:
+            # We can't do aiofiles easily for numpy, doing it sync for now
+            # as it's a small matrix and we use it as memory store
+            np.save(str(self._embeddings_path), self._centroids)
+
+    async def persist(self) -> None:
+        async with self._lock:
+            await self._persist_unsafe()
+
+    @property
+    def size(self) -> int:
+        return len(self._store)

segment_classifier/config.py

@@ -0,0 +1,53 @@
+from pydantic import BaseModel
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class ModelFeatureConfig(BaseModel):
+    """
+    Feature-based LLM model routing.
+
+    Selection priority (highest to lowest):
+    1. high_complexity  → use for ambiguous, deeply nested, multi-role segments
+    2. standard         → default for most unknown segments
+    3. fast             → simple segments with weak signals but not rule-matchable
+
+    Complexity is determined by:
+    - dom_depth > threshold
+    - child_tag_counts diversity (many unique tags = complex)
+    - text_density_ratio (very high or very low = complex)
+    - sibling_count == 0 (one-off sections = complex)
+    """
+    high_complexity_model: str = "anthropic/claude-opus-4"
+    standard_model: str = "anthropic/claude-sonnet-4-5"
+    fast_model: str = "anthropic/claude-haiku-4-5"
+
+    high_complexity_dom_depth_threshold: int = 6
+    high_complexity_unique_tag_threshold: int = 8
+    fast_model_max_dom_depth: int = 3
+
+
+class CacheConfig(BaseModel):
+    l1_cache_path: str = ".cache/l1_fingerprints.json"
+    l2_cache_path: str = ".cache/l2_clusters.json"
+    l2_embeddings_path: str = ".cache/l2_embeddings.npy"
+    l2_similarity_threshold: float = 0.85
+    l2_max_cluster_size: int = 50
+    persist_on_update: bool = True
+
+
+class ClassifierSettings(BaseSettings):
+    model_config = SettingsConfigDict(env_file=".env", env_prefix="CLASSIFIER_")
+
+    # LiteLLM
+    litellm_api_key: str = ""
+    litellm_batch_size: int = 20         # max segments per LLM batch call
+    litellm_max_concurrent_batches: int = 5
+    litellm_timeout_seconds: int = 60
+
+    # Pipeline
+    rule_based_confidence_threshold: float = 0.90
+    l1_min_confidence: float = 0.85
+    l2_min_confidence: float = 0.75
+
+    model_routing: ModelFeatureConfig = ModelFeatureConfig()
+    cache: CacheConfig = CacheConfig()

segment_classifier/models.py

@@ -0,0 +1,142 @@
+from enum import Enum
+from typing import Any
+from pydantic import BaseModel, Field
+
+
+class ClassificationStage(str, Enum):
+    RULE_BASED = "rule_based"
+    L1_EXACT_CACHE = "l1_exact_cache"
+    L2_FUZZY_CACHE = "l2_fuzzy_cache"
+    LLM = "llm"
+
+
+class SegmentPosition(str, Enum):
+    TOP = "top"           # top 5% of page
+    BOTTOM = "bottom"     # bottom 10% of page
+    MIDDLE = "middle"
+    UNKNOWN = "unknown"
+
+
+class ComponentType(str, Enum):
+    # Layout
+    LAYOUT_HEADER = "layout.header"
+    LAYOUT_FOOTER = "layout.footer"
+    LAYOUT_NAV = "layout.nav"
+    LAYOUT_SIDEBAR = "layout.sidebar"
+    LAYOUT_BREADCRUMB = "layout.breadcrumb"
+
+    # Collections
+    COLLECTION_PRODUCT_CARD = "collection.product_card"
+    COLLECTION_PRODUCT_LIST = "collection.product_list"
+    COLLECTION_BLOG_CARD = "collection.blog_card"
+    COLLECTION_BLOG_LIST = "collection.blog_list"
+    COLLECTION_NEWS_ITEM = "collection.news_item"
+    COLLECTION_NEWS_LIST = "collection.news_list"
+
+    # Sections
+    SECTION_HERO = "section.hero"
+    SECTION_FEATURE_GRID = "section.feature_grid"
+    SECTION_TESTIMONIAL = "section.testimonial"
+    SECTION_CTA = "section.cta"
+    SECTION_FAQ = "section.faq"
+    SECTION_PRICING = "section.pricing"
+
+    # UI Elements
+    UI_FORM = "ui.form"
+    UI_MODAL = "ui.modal"
+    UI_TABLE = "ui.table"
+    UI_CAROUSEL = "ui.carousel"
+    UI_PAGINATION = "ui.pagination"
+    UI_SEARCH = "ui.search"
+
+    # Content
+    CONTENT_ARTICLE = "content.article"
+    CONTENT_RICH_TEXT = "content.rich_text"
+    CONTENT_MEDIA = "content.media"
+
+    UNKNOWN = "unknown"
+
+
+class InputSegment(BaseModel):
+    """Raw segment from the page-segmenter tool."""
+    segment_id: str
+    page_url: str
+    page_slug: str
+    raw_html: str
+    text_content: str
+    position_hint: SegmentPosition = SegmentPosition.UNKNOWN
+    dom_position: str = ""                    # CSS selector path e.g. "main > section:nth-child(2)"
+    sibling_count: int = 0                    # how many same-fingerprint siblings on same page
+    url_path_segments: list[str] = Field(default_factory=list)  # e.g. ["products", "shoes"]
+
+
+class ClassifiedSegment(BaseModel):
+    """Segment with classification result and metadata."""
+    segment_id: str
+    page_url: str
+    page_slug: str
+    raw_html: str
+    text_content: str
+    position_hint: SegmentPosition
+
+    # Classification output
+    component_type: ComponentType
+    classification_stage: ClassificationStage
+    confidence: float = Field(ge=0.0, le=1.0)
+
+    # Fingerprint computed during pipeline
+    fingerprint_hash: str = ""
+    cluster_id: str | None = None
+
+    # LLM metadata (populated only for stage=LLM)
+    llm_model_used: str | None = None
+    llm_raw_response: str | None = None
+
+
+class FingerprintRecord(BaseModel):
+    """Stored in L1 cache: fingerprint → classification."""
+    fingerprint_hash: str
+    component_type: ComponentType
+    confidence: float
+    hit_count: int = 1
+    example_segment_id: str = ""
+
+
+class ClusterRecord(BaseModel):
+    """Stored in L2 cache: cluster of similar fingerprints."""
+    cluster_id: str
+    centroid_vector: list[float]
+    component_type: ComponentType
+    confidence: float
+    member_fingerprints: list[str] = Field(default_factory=list)
+
+
+class LLMClassificationRequest(BaseModel):
+    """Batch item sent to LLM."""
+    segment_id: str
+    fingerprint_hash: str
+    normalized_html: str          # skeleton only, no content
+    position_hint: SegmentPosition
+    sibling_count: int
+    url_hints: list[str]
+    dom_depth: int
+    child_tag_counts: dict[str, int]
+    text_density_ratio: float
+
+
+class LLMClassificationResult(BaseModel):
+    """Parsed result from LLM for one segment."""
+    segment_id: str
+    component_type: ComponentType
+    confidence: float
+    reasoning: str
+
+
+class PipelineResult(BaseModel):
+    """Final output of the full classification pipeline run."""
+    total_segments: int
+    classified: list[ClassifiedSegment]
+    stage_breakdown: dict[ClassificationStage, int]
+    llm_calls_made: int
+    llm_model_usage: dict[str, int]   # model_name → call count
+    cache_hit_rate: float

segment_classifier/pipeline.py

@@ -0,0 +1,173 @@
+import asyncio
+import logging
+from segment_classifier.models import (
+    InputSegment, ClassifiedSegment, PipelineResult, ClassificationStage, FingerprintRecord
+)
+from segment_classifier.config import ClassifierSettings
+from segment_classifier.utils.html_normalizer import NormalizedSegment
+from segment_classifier.stages.rule_based import RuleBasedClassifier
+from segment_classifier.stages.fingerprint import compute_fingerprint
+from segment_classifier.stages.fuzzy_cluster import FuzzyClusterStage
+from segment_classifier.stages.llm_classifier import LLMBatchClassifier
+from segment_classifier.cache.l1_cache import L1FingerprintCache
+from segment_classifier.cache.l2_cache import L2FuzzyCache
+
+logger = logging.getLogger(__name__)
+
+class ClassifierPipeline:
+    def __init__(self, settings: ClassifierSettings):
+        self.settings = settings
+        self.rule_classifier = RuleBasedClassifier(
+            confidence_threshold=settings.rule_based_confidence_threshold
+        )
+        self.l1_cache = L1FingerprintCache(
+            cache_path=settings.cache.l1_cache_path,
+            auto_persist_every=50
+        )
+        self.l2_cache = L2FuzzyCache(
+            cache_path=settings.cache.l2_cache_path,
+            embeddings_path=settings.cache.l2_embeddings_path,
+            similarity_threshold=settings.cache.l2_similarity_threshold,
+            max_cluster_size=settings.cache.l2_max_cluster_size,
+            persist_on_update=settings.cache.persist_on_update
+        )
+        self.fuzzy_stage = FuzzyClusterStage(self.l2_cache, settings.cache)
+        self.llm_classifier = LLMBatchClassifier(settings)
+
+    async def initialize(self) -> None:
+        """Load caches from disk."""
+        await asyncio.gather(
+            self.l1_cache.load(),
+            self.l2_cache.load()
+        )
+
+    async def shutdown(self) -> None:
+        """Persist caches to disk."""
+        await asyncio.gather(
+            self.l1_cache.persist(),
+            self.l2_cache.persist()
+        )
+
+    async def run(self, segments: list[InputSegment]) -> PipelineResult:
+        total_segments = len(segments)
+        classified: list[ClassifiedSegment] = []
+
+        # Precompute fingerprints concurrently
+        loop = asyncio.get_running_loop()
+
+        # We can run compute_fingerprint in an executor or direct async
+        # To avoid blocking event loop for too long if many segments, use executor
+        fingerprints_list = await asyncio.gather(
+            *[loop.run_in_executor(None, compute_fingerprint, seg) for seg in segments]
+        )
+
+        fingerprints: dict[str, tuple[NormalizedSegment, str]] = {
+            seg.segment_id: fp for seg, fp in zip(segments, fingerprints_list)
+        }
+
+        pending = segments.copy()
+
+        # Stage 1: Rule-based
+        next_pending = []
+        for segment in pending:
+            normalized, fp_hash = fingerprints[segment.segment_id]
+            result = self.rule_classifier.classify(segment, normalized)
+            if result:
+                classified.append(result)
+                await self.l1_cache.set(fp_hash, FingerprintRecord(
+                    fingerprint_hash=fp_hash,
+                    component_type=result.component_type,
+                    confidence=result.confidence,
+                    example_segment_id=segment.segment_id
+                ))
+            else:
+                next_pending.append(segment)
+        pending = next_pending
+
+        # Stage 2: L1 Exact Cache
+        next_pending = []
+        for segment in pending:
+            normalized, fp_hash = fingerprints[segment.segment_id]
+            record = await self.l1_cache.get(fp_hash)
+            if record and record.confidence >= self.settings.l1_min_confidence:
+                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=record.component_type,
+                    classification_stage=ClassificationStage.L1_EXACT_CACHE,
+                    confidence=record.confidence,
+                    fingerprint_hash=fp_hash
+                ))
+                await self.l1_cache.increment_hit(fp_hash)
+            else:
+                next_pending.append(segment)
+        pending = next_pending
+
+        # Stage 3: L2 Fuzzy Cache
+        next_pending = []
+        for segment in pending:
+            normalized, fp_hash = fingerprints[segment.segment_id]
+            result = await self.fuzzy_stage.classify(segment, normalized, fp_hash)
+            if result and result.confidence >= self.settings.l2_min_confidence:
+                classified.append(result)
+            else:
+                next_pending.append(segment)
+        pending = next_pending
+
+        # Stage 4: LLM Batch
+        llm_calls_made = 0
+        if pending:
+            llm_items = [
+                (seg, fingerprints[seg.segment_id][0], fingerprints[seg.segment_id][1])
+                for seg in pending
+            ]
+            llm_results = await self.llm_classifier.classify_batch(llm_items)
+
+            # For each LLM result, register in L1 + L2
+            for seg, result in zip(pending, llm_results):
+                normalized, fp_hash = fingerprints[seg.segment_id]
+                await self.l1_cache.set(fp_hash, FingerprintRecord(
+                    fingerprint_hash=fp_hash,
+                    component_type=result.component_type,
+                    confidence=result.confidence,
+                    example_segment_id=seg.segment_id
+                ))
+                await self.fuzzy_stage.register(
+                    fingerprint_hash=fp_hash,
+                    normalized=normalized,
+                    component_type=result.component_type,
+                    confidence=result.confidence
+                )
+
+            classified.extend(llm_results)
+
+            # Calculate total LLM batch calls
+            grouped_by_model: dict[str, int] = {}
+            for item in llm_items:
+                seg, norm, _ = item
+                model = self.llm_classifier.select_model(norm, seg)
+                grouped_by_model[model] = grouped_by_model.get(model, 0) + 1
+
+            for count in grouped_by_model.values():
+                import math
+                llm_calls_made += math.ceil(count / self.settings.litellm_batch_size)
+
+        stage_breakdown = {stage: 0 for stage in ClassificationStage}
+        for c in classified:
+            stage_breakdown[c.classification_stage] += 1
+
+        cache_hits = stage_breakdown[ClassificationStage.L1_EXACT_CACHE] + stage_breakdown[ClassificationStage.L2_FUZZY_CACHE]
+        cache_hit_rate = cache_hits / total_segments if total_segments > 0 else 0.0
+
+        return PipelineResult(
+            total_segments=total_segments,
+            classified=classified,
+            stage_breakdown=stage_breakdown,
+            llm_calls_made=llm_calls_made,
+            llm_model_usage=self.llm_classifier.model_usage,
+            cache_hit_rate=cache_hit_rate
+        )

segment_classifier/stages/__init__.py

@@ -0,0 +1,6 @@
+from .rule_based import RuleBasedClassifier
+from .fingerprint import compute_fingerprint
+from .fuzzy_cluster import FuzzyClusterStage
+from .llm_classifier import LLMBatchClassifier
+
+__all__ = ["RuleBasedClassifier", "compute_fingerprint", "FuzzyClusterStage", "LLMBatchClassifier"]

segment_classifier/stages/fingerprint.py

@@ -0,0 +1,10 @@
+from segment_classifier.models import InputSegment
+from segment_classifier.utils.html_normalizer import normalize_segment, NormalizedSegment
+
+def compute_fingerprint(segment: InputSegment) -> tuple[NormalizedSegment, str]:
+    """
+    Computes NormalizedSegment + fingerprint_hash for a given InputSegment.
+    """
+    normalized = normalize_segment(segment.raw_html, segment.text_content)
+    fingerprint_hash = normalized.fingerprint_hash()
+    return normalized, fingerprint_hash