docugraph 0.1.0__py3-none-any.whl

docugraph/ingestion/chunker.py

@@ -0,0 +1,382 @@
+"""Document chunking strategies."""
+
+import re
+from dataclasses import dataclass
+from enum import StrEnum
+
+from docugraph.core.models import Chunk, Document
+
+
+class ChunkingStrategy(StrEnum):
+    """Available chunking strategies."""
+
+    FIXED = "fixed"  # Fixed size with overlap
+    SEMANTIC = "semantic"  # By markdown headings/sections
+    SENTENCE = "sentence"  # By sentences
+    PARAGRAPH = "paragraph"  # By paragraphs
+
+
+@dataclass
+class ChunkerConfig:
+    """Configuration for the chunker."""
+
+    strategy: ChunkingStrategy = ChunkingStrategy.SEMANTIC
+    chunk_size: int = 1000  # Target characters per chunk
+    chunk_overlap: int = 200  # Overlap between chunks
+    min_chunk_size: int = 100  # Minimum chunk size
+    max_chunk_size: int = 2000  # Maximum chunk size
+
+
+class Chunker:
+    """Document chunker with multiple strategies."""
+
+    def __init__(self, config: ChunkerConfig | None = None):
+        """Initialize the chunker.
+
+        Args:
+            config: Chunker configuration.
+        """
+        self._config = config or ChunkerConfig()
+
+    def chunk_document(self, document: Document) -> list[Chunk]:
+        """Chunk a document into smaller pieces.
+
+        Args:
+            document: The document to chunk.
+
+        Returns:
+            List of chunks.
+        """
+        strategy = self._config.strategy
+
+        if strategy == ChunkingStrategy.SEMANTIC:
+            return self._chunk_semantic(document)
+        elif strategy == ChunkingStrategy.SENTENCE:
+            return self._chunk_sentences(document)
+        elif strategy == ChunkingStrategy.PARAGRAPH:
+            return self._chunk_paragraphs(document)
+        else:
+            return self._chunk_fixed(document)
+
+    def _chunk_fixed(self, document: Document) -> list[Chunk]:
+        """Chunk document with fixed size windows."""
+        content = document.content
+        chunks: list[Chunk] = []
+
+        chunk_size = self._config.chunk_size
+        overlap = self._config.chunk_overlap
+
+        start = 0
+        while start < len(content):
+            end = start + chunk_size
+
+            # Try to break at sentence/paragraph boundary
+            if end < len(content):
+                # Look for paragraph break
+                para_break = content.rfind("\n\n", start + chunk_size // 2, end + 100)
+                if para_break > start:
+                    end = para_break
+                else:
+                    # Look for sentence break
+                    sent_break = content.rfind(". ", start + chunk_size // 2, end + 50)
+                    if sent_break > start:
+                        end = sent_break + 1
+
+            chunk_content = content[start:end].strip()
+
+            if len(chunk_content) >= self._config.min_chunk_size:
+                chunks.append(
+                    Chunk(
+                        document_id=document.id,
+                        content=chunk_content,
+                        start_char=start,
+                        end_char=end,
+                        metadata={
+                            "source_url": document.source_url,
+                            "source_path": document.source_path,
+                            "title": document.title,
+                            "chunk_index": len(chunks),
+                        },
+                    )
+                )
+
+            start = end - overlap
+            if start >= len(content):
+                break
+
+        return chunks
+
+    def _chunk_semantic(self, document: Document) -> list[Chunk]:
+        """Chunk document by markdown sections."""
+        content = document.content
+        chunks: list[Chunk] = []
+
+        # Split by headings
+        heading_pattern = r"^(#{1,6})\s+(.+?)$"
+        sections: list[tuple[int, int, str, str]] = []  # (start, level, title, content)
+
+        lines = content.split("\n")
+        current_section_start = 0
+        current_title = document.title or "Introduction"
+        current_level = 0
+        section_lines: list[str] = []
+
+        for i, line in enumerate(lines):
+            match = re.match(heading_pattern, line)
+            if match:
+                # Save previous section
+                if section_lines:
+                    section_content = "\n".join(section_lines).strip()
+                    if section_content:
+                        sections.append(
+                            (
+                                current_section_start,
+                                current_level,
+                                current_title,
+                                section_content,
+                            )
+                        )
+
+                # Start new section
+                current_level = len(match.group(1))
+                current_title = match.group(2).strip()
+                current_section_start = sum(len(line) + 1 for line in lines[:i])
+                section_lines = [line]
+            else:
+                section_lines.append(line)
+
+        # Add last section
+        if section_lines:
+            section_content = "\n".join(section_lines).strip()
+            if section_content:
+                sections.append(
+                    (
+                        current_section_start,
+                        current_level,
+                        current_title,
+                        section_content,
+                    )
+                )
+
+        # Convert sections to chunks, splitting large ones
+        for start, level, title, section_content in sections:
+            if len(section_content) <= self._config.max_chunk_size:
+                if len(section_content) >= self._config.min_chunk_size:
+                    chunks.append(
+                        Chunk(
+                            document_id=document.id,
+                            content=section_content,
+                            start_char=start,
+                            end_char=start + len(section_content),
+                            metadata={
+                                "source_url": document.source_url,
+                                "source_path": document.source_path,
+                                "title": document.title,
+                                "section_title": title,
+                                "heading_level": level,
+                                "chunk_index": len(chunks),
+                            },
+                        )
+                    )
+            else:
+                # Split large sections with fixed chunking
+                sub_chunks = self._split_large_section(
+                    document, section_content, start, title, level, len(chunks)
+                )
+                chunks.extend(sub_chunks)
+
+        # If no chunks created (no headings found), fall back to fixed chunking
+        if not chunks:
+            return self._chunk_fixed(document)
+
+        return chunks
+
+    def _split_large_section(
+        self,
+        document: Document,
+        content: str,
+        base_start: int,
+        title: str,
+        level: int,
+        chunk_index_start: int,
+    ) -> list[Chunk]:
+        """Split a large section into smaller chunks."""
+        chunks: list[Chunk] = []
+
+        # Try to split by paragraphs first
+        paragraphs = content.split("\n\n")
+        current_chunk = ""
+        current_start = base_start
+        chunk_idx = chunk_index_start
+
+        for para in paragraphs:
+            if len(current_chunk) + len(para) + 2 <= self._config.chunk_size:
+                if current_chunk:
+                    current_chunk += "\n\n" + para
+                else:
+                    current_chunk = para
+            else:
+                # Save current chunk
+                if len(current_chunk) >= self._config.min_chunk_size:
+                    chunks.append(
+                        Chunk(
+                            document_id=document.id,
+                            content=current_chunk,
+                            start_char=current_start,
+                            end_char=current_start + len(current_chunk),
+                            metadata={
+                                "source_url": document.source_url,
+                                "source_path": document.source_path,
+                                "title": document.title,
+                                "section_title": title,
+                                "heading_level": level,
+                                "chunk_index": chunk_idx,
+                            },
+                        )
+                    )
+                    chunk_idx += 1
+
+                current_start = base_start + content.find(para)
+                current_chunk = para
+
+        # Add last chunk
+        if current_chunk and len(current_chunk) >= self._config.min_chunk_size:
+            chunks.append(
+                Chunk(
+                    document_id=document.id,
+                    content=current_chunk,
+                    start_char=current_start,
+                    end_char=current_start + len(current_chunk),
+                    metadata={
+                        "source_url": document.source_url,
+                        "source_path": document.source_path,
+                        "title": document.title,
+                        "section_title": title,
+                        "heading_level": level,
+                        "chunk_index": chunk_idx,
+                    },
+                )
+            )
+
+        return chunks
+
+    def _chunk_sentences(self, document: Document) -> list[Chunk]:
+        """Chunk document by sentences, grouping to target size."""
+        content = document.content
+        chunks: list[Chunk] = []
+
+        # Simple sentence splitting
+        sentence_pattern = r"(?<=[.!?])\s+"
+        sentences = re.split(sentence_pattern, content)
+
+        current_chunk = ""
+        current_start = 0
+
+        for sentence in sentences:
+            if len(current_chunk) + len(sentence) + 1 <= self._config.chunk_size:
+                if current_chunk:
+                    current_chunk += " " + sentence
+                else:
+                    current_chunk = sentence
+            else:
+                if len(current_chunk) >= self._config.min_chunk_size:
+                    chunks.append(
+                        Chunk(
+                            document_id=document.id,
+                            content=current_chunk,
+                            start_char=current_start,
+                            end_char=current_start + len(current_chunk),
+                            metadata={
+                                "source_url": document.source_url,
+                                "source_path": document.source_path,
+                                "title": document.title,
+                                "chunk_index": len(chunks),
+                            },
+                        )
+                    )
+
+                current_start = content.find(sentence, current_start)
+                current_chunk = sentence
+
+        # Add last chunk
+        if current_chunk and len(current_chunk) >= self._config.min_chunk_size:
+            chunks.append(
+                Chunk(
+                    document_id=document.id,
+                    content=current_chunk,
+                    start_char=current_start,
+                    end_char=current_start + len(current_chunk),
+                    metadata={
+                        "source_url": document.source_url,
+                        "source_path": document.source_path,
+                        "title": document.title,
+                        "chunk_index": len(chunks),
+                    },
+                )
+            )
+
+        return chunks
+
+    def _chunk_paragraphs(self, document: Document) -> list[Chunk]:
+        """Chunk document by paragraphs, grouping to target size."""
+        content = document.content
+        chunks: list[Chunk] = []
+
+        paragraphs = content.split("\n\n")
+
+        current_chunk = ""
+        current_start = 0
+
+        for para in paragraphs:
+            para = para.strip()
+            if not para:
+                continue
+
+            if len(current_chunk) + len(para) + 2 <= self._config.chunk_size:
+                if current_chunk:
+                    current_chunk += "\n\n" + para
+                else:
+                    current_chunk = para
+            else:
+                if len(current_chunk) >= self._config.min_chunk_size:
+                    chunks.append(
+                        Chunk(
+                            document_id=document.id,
+                            content=current_chunk,
+                            start_char=current_start,
+                            end_char=current_start + len(current_chunk),
+                            metadata={
+                                "source_url": document.source_url,
+                                "source_path": document.source_path,
+                                "title": document.title,
+                                "chunk_index": len(chunks),
+                            },
+                        )
+                    )
+
+                current_start = content.find(para, current_start)
+                current_chunk = para
+
+        # Add last chunk
+        if current_chunk and len(current_chunk) >= self._config.min_chunk_size:
+            chunks.append(
+                Chunk(
+                    document_id=document.id,
+                    content=current_chunk,
+                    start_char=current_start,
+                    end_char=current_start + len(current_chunk),
+                    metadata={
+                        "source_url": document.source_url,
+                        "source_path": document.source_path,
+                        "title": document.title,
+                        "chunk_index": len(chunks),
+                    },
+                )
+            )
+
+        return chunks
+
+
+def get_chunker(config: ChunkerConfig | None = None) -> Chunker:
+    """Get a chunker instance."""
+    return Chunker(config)

docugraph/ingestion/crawler.py

@@ -0,0 +1,305 @@
+"""Web crawler using Crawl4AI."""
+
+import asyncio
+import hashlib
+from datetime import UTC, datetime
+from urllib.parse import urljoin, urlparse
+
+from docugraph.core.config import CrawlerConfig, get_config
+from docugraph.core.models import ContentType, Document
+
+
+class DocCrawler:
+    """Web crawler for documentation sites using Crawl4AI."""
+
+    def __init__(self, config: CrawlerConfig | None = None):
+        """Initialize the crawler.
+
+        Args:
+            config: Crawler configuration. If None, uses default from config.
+        """
+        self._config = config or get_config().crawler
+        self._cache_dir = get_config().storage.data_dir / "crawl_cache"
+        self._cache_dir.mkdir(parents=True, exist_ok=True)
+
+    def _get_cache_key(self, url: str) -> str:
+        """Generate a cache key for a URL."""
+        return hashlib.sha256(url.encode()).hexdigest()[:16]
+
+    def _get_cached(self, url: str) -> Document | None:
+        """Get cached document if available and fresh."""
+        cache_key = self._get_cache_key(url)
+        cache_file = self._cache_dir / f"{cache_key}.json"
+
+        if not cache_file.exists():
+            return None
+
+        import json
+
+        try:
+            with open(cache_file) as f:
+                data = json.load(f)
+
+            # Check TTL. A cache entry written before timestamps carried a
+            # timezone parses as naive, so it is read as UTC rather than
+            # subtracted from an aware "now" (which raises).
+            cached_at = datetime.fromisoformat(data["cached_at"])
+            if cached_at.tzinfo is None:
+                cached_at = cached_at.replace(tzinfo=UTC)
+            age = (datetime.now(UTC) - cached_at).total_seconds()
+            if age > self._config.cache_ttl:
+                return None
+
+            return Document(**data["document"])
+        except Exception:
+            return None
+
+    def _save_cache(self, url: str, document: Document) -> None:
+        """Save document to cache."""
+        cache_key = self._get_cache_key(url)
+        cache_file = self._cache_dir / f"{cache_key}.json"
+
+        import json
+
+        data = {
+            "cached_at": datetime.now(UTC).isoformat(),
+            "document": document.model_dump(mode="json"),
+        }
+
+        with open(cache_file, "w") as f:
+            json.dump(data, f)
+
+    async def crawl_single(
+        self,
+        url: str,
+        use_cache: bool = True,
+    ) -> Document:
+        """Crawl a single URL.
+
+        Args:
+            url: The URL to crawl.
+            use_cache: Whether to use cached results.
+
+        Returns:
+            The crawled document.
+        """
+        # Check cache first
+        if use_cache:
+            cached = self._get_cached(url)
+            if cached:
+                return cached
+
+        from crawl4ai import AsyncWebCrawler, BrowserConfig, CrawlerRunConfig
+        from crawl4ai.content_filter_strategy import PruningContentFilter
+        from crawl4ai.markdown_generation_strategy import DefaultMarkdownGenerator
+
+        browser_config = BrowserConfig(
+            headless=True,
+            verbose=False,
+        )
+
+        crawler_config = CrawlerRunConfig(
+            markdown_generator=DefaultMarkdownGenerator(
+                content_filter=PruningContentFilter(
+                    threshold=0.48,
+                    threshold_type="fixed",
+                    min_word_threshold=0,
+                )
+            ),
+            wait_until="domcontentloaded",
+        )
+
+        async with AsyncWebCrawler(config=browser_config) as crawler:
+            result = await crawler.arun(url=url, config=crawler_config)
+
+            if not result.success:
+                raise RuntimeError(f"Failed to crawl {url}: {result.error_message}")
+
+            # Extract title from metadata or content
+            title = result.metadata.get("title") if result.metadata else None
+            if not title and result.markdown:
+                # Try to extract from first heading
+                lines = result.markdown.split("\n")
+                for line in lines:
+                    if line.startswith("# "):
+                        title = line[2:].strip()
+                        break
+
+            document = Document(
+                source_url=url,
+                title=title,
+                content=result.markdown.fit_markdown
+                if hasattr(result.markdown, "fit_markdown")
+                else result.markdown,
+                content_type=ContentType.MARKDOWN,
+                metadata={
+                    "raw_html_length": len(result.html) if result.html else 0,
+                    "links_count": len(result.links.get("internal", [])) if result.links else 0,
+                    "crawled_at": datetime.now(UTC).isoformat(),
+                },
+            )
+
+            # Cache the result
+            if use_cache:
+                self._save_cache(url, document)
+
+            return document
+
+    async def crawl_site(
+        self,
+        start_url: str,
+        max_pages: int = 100,
+        url_pattern: str | None = None,
+        use_cache: bool = True,
+    ) -> list[Document]:
+        """Crawl a documentation site starting from a URL.
+
+        Args:
+            start_url: The starting URL.
+            max_pages: Maximum number of pages to crawl.
+            url_pattern: Optional regex pattern to filter URLs.
+            use_cache: Whether to use cached results.
+
+        Returns:
+            List of crawled documents.
+        """
+        from crawl4ai import AsyncWebCrawler, BrowserConfig, CrawlerRunConfig
+        from crawl4ai.content_filter_strategy import PruningContentFilter
+        from crawl4ai.markdown_generation_strategy import DefaultMarkdownGenerator
+
+        browser_config = BrowserConfig(
+            headless=True,
+            verbose=False,
+        )
+
+        crawler_config = CrawlerRunConfig(
+            markdown_generator=DefaultMarkdownGenerator(
+                content_filter=PruningContentFilter(
+                    threshold=0.48,
+                    threshold_type="fixed",
+                    min_word_threshold=0,
+                )
+            ),
+            wait_until="domcontentloaded",
+        )
+
+        # Parse base domain for filtering
+        parsed_start = urlparse(start_url)
+        base_domain = parsed_start.netloc
+        parsed_start.path.rsplit("/", 1)[0] if "/" in parsed_start.path else ""
+
+        visited: set[str] = set()
+        to_visit: list[str] = [start_url]
+        documents: list[Document] = []
+
+        import re
+
+        url_regex = re.compile(url_pattern) if url_pattern else None
+
+        async with AsyncWebCrawler(config=browser_config) as crawler:
+            while to_visit and len(documents) < max_pages:
+                # Get next URL to crawl
+                url = to_visit.pop(0)
+
+                if url in visited:
+                    continue
+
+                visited.add(url)
+
+                # Check cache
+                if use_cache:
+                    cached = self._get_cached(url)
+                    if cached:
+                        documents.append(cached)
+                        continue
+
+                try:
+                    result = await crawler.arun(url=url, config=crawler_config)
+
+                    if not result.success:
+                        continue
+
+                    # Extract title
+                    title = result.metadata.get("title") if result.metadata else None
+                    if not title and result.markdown:
+                        lines = result.markdown.split("\n")
+                        for line in lines:
+                            if line.startswith("# "):
+                                title = line[2:].strip()
+                                break
+
+                    content = (
+                        result.markdown.fit_markdown
+                        if hasattr(result.markdown, "fit_markdown")
+                        else result.markdown
+                    )
+
+                    document = Document(
+                        source_url=url,
+                        title=title,
+                        content=content,
+                        content_type=ContentType.MARKDOWN,
+                        metadata={
+                            "raw_html_length": len(result.html) if result.html else 0,
+                            "crawled_at": datetime.now(UTC).isoformat(),
+                        },
+                    )
+
+                    documents.append(document)
+
+                    if use_cache:
+                        self._save_cache(url, document)
+
+                    # Extract and queue internal links
+                    if result.links and "internal" in result.links:
+                        for link in result.links["internal"]:
+                            href = link.get("href", "")
+                            if not href:
+                                continue
+
+                            # Resolve relative URLs
+                            full_url = urljoin(url, href)
+                            parsed = urlparse(full_url)
+
+                            # Filter to same domain
+                            if parsed.netloc != base_domain:
+                                continue
+
+                            # Filter by pattern if provided
+                            if url_regex and not url_regex.search(full_url):
+                                continue
+
+                            # Skip anchors and query params
+                            clean_url = f"{parsed.scheme}://{parsed.netloc}{parsed.path}"
+
+                            if clean_url not in visited and clean_url not in to_visit:
+                                to_visit.append(clean_url)
+
+                except Exception as e:
+                    # Log error but continue crawling
+                    print(f"Error crawling {url}: {e}")
+                    continue
+
+                # Respect rate limit
+                await asyncio.sleep(1.0 / self._config.rate_limit)
+
+        return documents
+
+    def crawl_single_sync(self, url: str, use_cache: bool = True) -> Document:
+        """Synchronous wrapper for crawl_single."""
+        return asyncio.run(self.crawl_single(url, use_cache))
+
+    def crawl_site_sync(
+        self,
+        start_url: str,
+        max_pages: int = 100,
+        url_pattern: str | None = None,
+        use_cache: bool = True,
+    ) -> list[Document]:
+        """Synchronous wrapper for crawl_site."""
+        return asyncio.run(self.crawl_site(start_url, max_pages, url_pattern, use_cache))
+
+
+def get_crawler() -> DocCrawler:
+    """Get the default crawler instance."""
+    return DocCrawler()