sec2md 0.1.0__py3-none-any.whl

sec2md/chunker/markdown_blocks.py

@@ -0,0 +1,116 @@
+import re
+from abc import ABC
+from typing import List
+
+
+def estimate_tokens(text: str) -> int:
+    """
+    Estimate token count using character/4 heuristic.
+
+    This is a simple approximation. For exact token counting,
+    use your embedding provider's tokenizer.
+    """
+    return max(1, len(text) // 4)
+
+
+def split_sentences(text: str) -> List[str]:
+    """Simple regex-based sentence splitter"""
+    # Split on .!? followed by whitespace and capital letter or end of string
+    # Handles common abbreviations like Mr., Dr., Inc., etc.
+    sentences = re.split(r'(?<=[.!?])\s+(?=[A-Z])', text)
+    return [s.strip() for s in sentences if s.strip()]
+
+
+class BaseBlock(ABC):
+    block_type: str
+
+    def __init__(self, content: str, page: int):
+        self.content = content
+        self.page = page
+
+    @property
+    def tokens(self) -> int:
+        return estimate_tokens(self.content)
+
+
+class Sentence:
+
+    def __init__(self, content: str):
+        self.content = content
+
+    @property
+    def tokens(self) -> int:
+        return estimate_tokens(self.content)
+
+
+class TextBlock(BaseBlock):
+    block_type: str = 'Text'
+
+    def __init__(self, content: str, page: int):
+        super().__init__(content=content, page=page)
+
+    @property
+    def sentences(self) -> List[Sentence]:
+        """Returns the text block sentences"""
+        return [Sentence(content=content) for content in split_sentences(self.content)]
+
+    @classmethod
+    def from_sentences(cls, sentences: List[Sentence], page: int):
+        content = " ".join([sentence.content for sentence in sentences])
+        return cls(content=content, page=page)
+
+
+class AudioParagraphBlock(BaseBlock):
+    block_type: str = "Text"
+
+    def __init__(self, content: str, page: int, paragraph_id: int, audio_start: float, audio_end: float):
+        super().__init__(content=content, page=page)
+        self.paragraph_id = paragraph_id
+        self.audio_start = audio_start
+        self.audio_end = audio_end
+
+    @property
+    def sentences(self) -> List[Sentence]:
+        """Returns the text block sentences"""
+        return [Sentence(content=content) for content in split_sentences(self.content)]
+
+    def format(self) -> dict:
+        """Formats the audio paragraphs"""
+        return {"id": self.paragraph_id, "content": self.content, "start": self.audio_start, "end": self.audio_end}
+
+
+class TableBlock(BaseBlock):
+    block_type: str = 'Table'
+
+    def __init__(self, content: str, page: int):
+        super().__init__(content=content, page=page)
+        self.content = self._to_minified_markdown()
+
+    def _to_minified_markdown(self) -> str:
+        """Returns the table in a Minified Markdown format"""
+        lines = self.content.split('\n')
+        cleaned_lines = []
+
+        for i, line in enumerate(lines):
+            if not line.strip():
+                continue
+
+            parts = line.split('|')
+            cleaned_parts = [re.sub(r'\s+', ' ', part.strip()) for part in parts]
+            cleaned_line = '|'.join(cleaned_parts)
+
+            if i == 1:
+                num_cols = len(cleaned_parts) - 1
+                separator = '|' + '|'.join(['---'] * num_cols) + '|'
+                cleaned_lines.append(separator)
+            else:
+                cleaned_lines.append(cleaned_line)
+
+        return '\n'.join(cleaned_lines)
+
+
+class HeaderBlock(BaseBlock):
+    block_type = 'Header'
+
+    def __init__(self, content: str, page: int):
+        super().__init__(content=content, page=page)

sec2md/chunker/markdown_chunk.py

@@ -0,0 +1,76 @@
+from typing import List, Optional
+
+from sec2md.chunker.markdown_blocks import BaseBlock
+
+
+class MarkdownChunk:
+    """Represents a chunk of markdown content that can be embedded"""
+
+    def __init__(self, blocks: List[BaseBlock], header: Optional[str] = None):
+        """Initialize a markdown chunk with blocks and optional header for embedding"""
+        self.vector: Optional[List[float]] = None
+        self.blocks = blocks
+        self.page = blocks[0].page
+        self.header = header
+
+    def set_vector(self, vector: List[float]):
+        """Set the vector embedding for this chunk"""
+        self.vector = vector
+
+    @property
+    def content(self) -> str:
+        """Get the text content of this chunk"""
+        return "\n".join([block.content for block in self.blocks])
+
+    @property
+    def data(self) -> List[dict]:
+        """Returns a list of block data grouped by page with ONLY the chunk's content"""
+        page_blocks = {}
+
+        for block in self.blocks:
+            if block.page not in page_blocks:
+                page_blocks[block.page] = []
+            page_blocks[block.page].append(block)
+
+        page_content_data = []
+        for page, blocks in page_blocks.items():
+            # Only include the content from blocks in THIS chunk, not full page content
+            page_content = "\n".join(block.content for block in blocks)
+            if not page_content.strip():
+                continue
+
+            page_content_data.append({
+                "page": page,
+                "content": page_content
+            })
+
+        return sorted(page_content_data, key=lambda x: x["page"])
+
+    @property
+    def pages(self) -> List[dict]:
+        """Returns a list of pages with ONLY this chunk's content (not full page content)"""
+        return self.data
+
+    @property
+    def embedding_text(self) -> str:
+        """Get the text to use for embedding, with optional header prepended"""
+        if self.header:
+            return f"{self.header}\n\n...\n\n{self.content}"
+        return self.content
+
+    @property
+    def has_table(self) -> bool:
+        """Returns True if this chunk contains one or more table blocks"""
+        return any(block.block_type == 'Table' for block in self.blocks)
+
+    @property
+    def num_tokens(self) -> int:
+        """Returns the total number of tokens in this chunk"""
+        return sum(block.tokens for block in self.blocks)
+
+    def __repr__(self):
+        return f"MarkdownChunk(page={self.page}, blocks={len(self.blocks)})"
+
+    def _repr_markdown_(self):
+        """This method is called by IPython to display as Markdown"""
+        return self.content

sec2md/chunker/markdown_chunker.py

@@ -0,0 +1,234 @@
+import logging
+from typing import Union, Tuple, List, Dict
+
+from sec2md.chunker.markdown_chunk import MarkdownChunk
+from sec2md.chunker.markdown_blocks import BaseBlock, TextBlock, TableBlock, HeaderBlock
+
+logger = logging.getLogger(__name__)
+
+
+class MarkdownChunker:
+    """Splits markdown content into chunks"""
+
+    def __init__(self, chunk_size: int = 512, chunk_overlap: int = 128):
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+
+    def split(self, pages: List[Dict[str, Union[int, str]]], header: str = None) -> List[MarkdownChunk]:
+        """Split the pages into chunks with optional header for embedding context"""
+        blocks = self._split_into_blocks(pages=pages)
+        return self._chunk_blocks(blocks=blocks, header=header)
+
+    def chunk_text(self, text: str) -> List[str]:
+        """Chunk a single text string into multiple chunks"""
+        pages = [{"page": 0, "content": text}]
+        chunks = self.split(pages=pages)
+        return [chunk.content for chunk in chunks]
+
+    @staticmethod
+    def _split_into_blocks(pages: List[Dict[str, Union[int, str]]]):
+        """Splits the page into blocks"""
+        blocks = []
+        table_content = ""
+        last_page = None
+
+        for page in pages:
+            last_page = page['page']
+            for line in page['content'].split('\n'):
+                if table_content and not MarkdownChunker._is_table_line(line):
+                    block = TableBlock(content=table_content, page=page['page'])
+                    blocks.append(block)
+                    table_content = ""
+
+                if line.startswith("#"):
+                    block = HeaderBlock(content=line, page=page['page'])
+                    blocks.append(block)
+
+                elif MarkdownChunker._is_table_line(line):
+                    table_content += f"{line}\n"
+
+                else:
+                    block = TextBlock(content=line, page=page['page'])
+                    blocks.append(block)
+
+        if table_content and last_page is not None:
+            block = TableBlock(content=table_content, page=last_page)
+            blocks.append(block)
+
+        return blocks
+
+    @staticmethod
+    def _is_table_line(line: str) -> bool:
+        import re
+        if '|' not in line:
+            return False
+        stripped = line.strip()
+        if not stripped:
+            return False
+        align_pattern = re.compile(r'^\s*:?-+:?\s*$')
+        cells = [c.strip() for c in stripped.strip('|').split('|')]
+        if all(align_pattern.match(c) for c in cells):
+            return True
+        return True
+
+    def _chunk_blocks(self, blocks: List[BaseBlock], header: str = None) -> List[MarkdownChunk]:
+        """Converts the blocks to chunks"""
+        chunks = []
+        chunk_blocks = []
+        num_tokens = 0
+
+        for i, block in enumerate(blocks):
+            next_block = blocks[i + 1] if i + 1 < len(blocks) else None
+
+            if block.block_type == 'Text':
+                chunk_blocks, num_tokens, chunks = self._process_text_block(
+                    block, chunk_blocks, num_tokens, chunks, header
+                )
+
+            elif block.block_type == 'Table':
+                chunk_blocks, num_tokens, chunks = self._process_table_block(
+                    block, chunk_blocks, num_tokens, chunks, blocks, i, header
+                )
+
+            else:
+                chunk_blocks, num_tokens, chunks = self._process_header_table_block(
+                    block, chunk_blocks, num_tokens, chunks, next_block, header
+                )
+
+        if chunk_blocks:
+            chunks.append(MarkdownChunk(blocks=chunk_blocks, header=header))
+
+        return chunks
+
+    def _process_text_block(self, block: TextBlock, chunk_blocks: List[BaseBlock], num_tokens: int,
+                            chunks: List[MarkdownChunk], header: str = None):
+        """Process a text block by breaking it into sentences if needed"""
+        sentences = []
+        sentences_tokens = 0
+
+        for sentence in block.sentences:
+            if num_tokens + sentences_tokens + sentence.tokens > self.chunk_size:
+                if sentences:
+                    new_block = TextBlock.from_sentences(sentences=sentences, page=block.page)
+                    chunk_blocks.append(new_block)
+                    num_tokens += sentences_tokens
+
+                chunks, chunk_blocks, num_tokens = self._create_chunk(chunks=chunks, blocks=chunk_blocks, header=header)
+
+                sentences = [sentence]
+                sentences_tokens = sentence.tokens
+
+            else:
+                sentences.append(sentence)
+                sentences_tokens += sentence.tokens
+
+        if sentences:
+            new_block = TextBlock.from_sentences(sentences=sentences, page=block.page)
+            chunk_blocks.append(new_block)
+            num_tokens += sentences_tokens
+
+        return chunk_blocks, num_tokens, chunks
+
+    def _process_table_block(self, block: BaseBlock, chunk_blocks: List[BaseBlock], num_tokens: int,
+                             chunks: List[MarkdownChunk], all_blocks: List[BaseBlock], block_idx: int, header: str = None):
+        """Process a table block with optional header backtrack"""
+        context = []
+        context_tokens = 0
+
+        # Backtrack for header only if 1-2 short blocks precede
+        count = 0
+        for j in range(block_idx - 1, -1, -1):
+            prev = all_blocks[j]
+            if prev.page != block.page:
+                break
+            if prev.block_type == 'Header':
+                if context_tokens + prev.tokens <= 128:
+                    context.insert(0, prev)
+                    context_tokens += prev.tokens
+                break
+            elif prev.block_type == 'Text' and prev.content.strip():
+                count += 1
+                if count > 2:
+                    break
+                if context_tokens + prev.tokens <= 128:
+                    context.insert(0, prev)
+                    context_tokens += prev.tokens
+                else:
+                    break
+
+        if num_tokens + context_tokens + block.tokens > self.chunk_size:
+            if chunk_blocks:
+                chunks, chunk_blocks, num_tokens = self._create_chunk(chunks=chunks, blocks=chunk_blocks, header=header)
+
+            # If we're backtracking context and the last chunk is ONLY that context, remove it
+            if context and chunks and len(chunks[-1].blocks) == len(context):
+                if all(chunks[-1].blocks[i] == context[i] for i in range(len(context))):
+                    chunks.pop()
+
+            chunk_blocks = context + [block]
+            num_tokens = context_tokens + block.tokens
+        else:
+            chunk_blocks.extend(context + [block])
+            num_tokens += context_tokens + block.tokens
+
+        return chunk_blocks, num_tokens, chunks
+
+    def _process_header_table_block(self, block: BaseBlock, chunk_blocks: List[BaseBlock], num_tokens: int,
+                                    chunks: List[MarkdownChunk], next_block: BaseBlock, header: str = None):
+        """Process a header block"""
+        if not chunk_blocks:
+            chunk_blocks.append(block)
+            num_tokens += block.tokens
+            return chunk_blocks, num_tokens, chunks
+
+        # Don't split if current content is small and next is a table
+        if next_block and next_block.block_type == 'Table' and num_tokens < self.chunk_overlap:
+            chunk_blocks.append(block)
+            num_tokens += block.tokens
+            return chunk_blocks, num_tokens, chunks
+
+        if num_tokens + block.tokens > self.chunk_size:
+            chunks, chunk_blocks, num_tokens = self._create_chunk(chunks=chunks, blocks=chunk_blocks, header=header)
+            chunk_blocks.append(block)
+            num_tokens += block.tokens
+        else:
+            chunk_blocks.append(block)
+            num_tokens += block.tokens
+
+        return chunk_blocks, num_tokens, chunks
+
+    def _create_chunk(self, chunks: List[MarkdownChunk], blocks: List[BaseBlock], header: str = None) -> Tuple[
+        List[MarkdownChunk], List[BaseBlock], int]:
+        """Creates a chunk, and return a new list of blocks that """
+        chunks.append(MarkdownChunk(blocks=blocks, header=header))
+
+        if not self.chunk_overlap:
+            return chunks, [], 0
+
+        overlap_tokens = 0
+        overlap_blocks = []
+
+        for block in reversed(blocks):
+            if block.block_type == "Text":
+                sentences = []
+
+                for sentence in reversed(block.sentences):
+
+                    if overlap_tokens + sentence.tokens > self.chunk_overlap:
+                        text_block = TextBlock.from_sentences(sentences=sentences, page=block.page)
+                        overlap_blocks.insert(0, text_block)
+                        return chunks, overlap_blocks, overlap_tokens
+
+                    else:
+                        sentences.insert(0, sentence)
+                        overlap_tokens += sentence.tokens
+
+            else:
+                if overlap_tokens + block.tokens > self.chunk_overlap:
+                    return chunks, overlap_blocks, overlap_tokens
+
+                else:
+                    overlap_blocks.insert(0, block)
+                    overlap_tokens += block.tokens
+
+        return chunks, [], 0

sec2md/chunking.py

@@ -0,0 +1,66 @@
+"""Chunking utilities for page-aware markdown splitting."""
+
+from typing import List, Optional
+from sec2md.models import Page, Section
+from sec2md.chunker.markdown_chunker import MarkdownChunker
+from sec2md.chunker.markdown_chunk import MarkdownChunk
+
+
+def chunk_pages(
+    pages: List[Page],
+    chunk_size: int = 512,
+    chunk_overlap: int = 128,
+    header: Optional[str] = None
+) -> List[MarkdownChunk]:
+    """
+    Chunk pages into overlapping markdown chunks.
+
+    Args:
+        pages: List of Page objects
+        chunk_size: Target chunk size in tokens (estimated as chars/4)
+        chunk_overlap: Overlap between chunks in tokens
+        header: Optional header to prepend to each chunk's embedding_text
+
+    Returns:
+        List of MarkdownChunk objects with page tracking
+
+    Example:
+        >>> pages = sec2md.convert_to_markdown(html, return_pages=True)
+        >>> chunks = sec2md.chunk_pages(pages, chunk_size=512)
+        >>> for chunk in chunks:
+        ...     print(f"Page {chunk.page}: {chunk.content[:100]}...")
+    """
+    chunker = MarkdownChunker(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    pages_data = [{"page": p.number, "content": p.content} for p in pages]
+    return chunker.split(pages=pages_data, header=header)
+
+
+def chunk_section(
+    section: Section,
+    chunk_size: int = 512,
+    chunk_overlap: int = 128,
+    header: Optional[str] = None
+) -> List[MarkdownChunk]:
+    """
+    Chunk a filing section into overlapping markdown chunks.
+
+    Args:
+        section: Section object from extract_sections()
+        chunk_size: Target chunk size in tokens (estimated as chars/4)
+        chunk_overlap: Overlap between chunks in tokens
+        header: Optional header to prepend to each chunk's embedding_text
+
+    Returns:
+        List of MarkdownChunk objects
+
+    Example:
+        >>> sections = sec2md.extract_sections(pages, filing_type="10-K")
+        >>> risk = sec2md.get_section(sections, Item10K.RISK_FACTORS)
+        >>> chunks = sec2md.chunk_section(risk, chunk_size=512)
+    """
+    return chunk_pages(
+        pages=section.pages,
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        header=header
+    )

sec2md/core.py

@@ -0,0 +1,93 @@
+"""Core conversion functionality."""
+
+from typing import overload, List
+from sec2md.utils import is_url, fetch
+from sec2md.parser import Parser
+from sec2md.models import Page
+
+
+@overload
+def convert_to_markdown(
+    source: str | bytes,
+    *,
+    user_agent: str | None = None,
+    return_pages: bool = False,
+) -> str: ...
+
+
+@overload
+def convert_to_markdown(
+    source: str | bytes,
+    *,
+    user_agent: str | None = None,
+    return_pages: bool = True,
+) -> List[Page]: ...
+
+
+def convert_to_markdown(
+    source: str | bytes,
+    *,
+    user_agent: str | None = None,
+    return_pages: bool = False,
+) -> str | List[Page]:
+    """
+    Convert SEC filing HTML to Markdown.
+
+    Args:
+        source: URL or HTML string/bytes
+        user_agent: User agent for EDGAR requests (required for sec.gov URLs)
+        return_pages: If True, returns List[Page] instead of markdown string
+
+    Returns:
+        Markdown string (default) or List[Page] if return_pages=True
+
+    Raises:
+        ValueError: If source appears to be PDF content or other non-HTML format
+
+    Examples:
+        >>> # From URL - get markdown
+        >>> md = convert_to_markdown(
+        ...     "https://www.sec.gov/Archives/edgar/data/.../10k.htm",
+        ...     user_agent="Lucas Astorian <lucas@intellifin.ai>"
+        ... )
+
+        >>> # Get pages for section extraction
+        >>> pages = convert_to_markdown(filing.html(), return_pages=True)
+
+        >>> # With edgartools
+        >>> from edgar import Company, set_identity
+        >>> set_identity("Lucas Astorian <lucas@intellifin.ai>")
+        >>> company = Company('AAPL')
+        >>> filing = company.get_filings(form="10-K").latest()
+        >>> md = convert_to_markdown(filing.html())
+    """
+    # Handle bytes input
+    if isinstance(source, bytes):
+        # Check if it's PDF
+        if source.startswith(b'%PDF'):
+            raise ValueError(
+                "PDF content detected. This library only supports HTML input. "
+                "Please extract HTML from the filing first."
+            )
+        source = source.decode('utf-8', errors='ignore')
+
+    # Check for PDF in string
+    if isinstance(source, str) and source.strip().startswith('%PDF'):
+        raise ValueError(
+            "PDF content detected. This library only supports HTML input. "
+            "Please extract HTML from the filing first."
+        )
+
+    # Fetch from URL if needed
+    if is_url(source):
+        html = fetch(source, user_agent=user_agent)
+    else:
+        html = source
+
+    # Parse and convert
+    parser = Parser(html)
+
+    if return_pages:
+        return parser.get_pages()
+    else:
+        return parser.markdown()