agentic_rag/services/knowledge/processors/text_processor.py

"""Text modal processor — chunking, entity extraction, and annotation."""

from agentic_rag.services.knowledge.content_list import ContentItem, ContentType
from agentic_rag.services.knowledge.processors.base import BaseModalProcessor


class TextModalProcessor(BaseModalProcessor):
    """Process text content — chunking, cleaning, entity extraction."""

    content_type = ContentType.TEXT
    description = "Clean and chunk text content."

    def __init__(self, chunk_size: int = 512, chunk_overlap: int = 50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    async def process(self, item: ContentItem) -> ContentItem:
        """Clean and optionally chunk text content."""
        if item.type != ContentType.TEXT:
            return item

        # Basic cleaning
        text = item.text.strip()
        text = self._clean_text(text)

        if not text:
            return item

        item.text = text
        return item

    async def process_batch(self, items: list[ContentItem]) -> list[ContentItem]:
        """Process text items. For texts longer than chunk_size, split into chunks."""
        results = []
        for item in items:
            processed = await self.process(item)
            if processed.text and len(processed.text) > self.chunk_size:
                # Split long text into overlapping chunks
                chunks = self._chunk_text(processed.text)
                for i, chunk in enumerate(chunks):
                    chunk_item = ContentItem(
                        type=ContentType.TEXT,
                        text=chunk,
                        page_idx=processed.page_idx,
                        metadata={
                            **processed.metadata,
                            "chunk_index": i,
                            "chunk_count": len(chunks),
                            "parent_item": id(processed),
                        },
                    )
                    results.append(chunk_item)
            else:
                results.append(processed)
        return results

    def _clean_text(self, text: str) -> str:
        """Clean up text — normalize whitespace, remove artifacts."""
        import re
        # Collapse multiple newlines
        text = re.sub(r'\n{3,}', '\n\n', text)
        # Collapse multiple spaces
        text = re.sub(r' {3,}', '  ', text)
        # Remove control characters except newlines
        text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
        return text.strip()

    def _chunk_text(self, text: str) -> list[str]:
        """Split text into chunks at Chinese/English sentence boundaries.

        Overlap is sentence-aware — the last complete sentence(s) from the
        previous chunk are prepended to the next one, so no chunk starts
        with a mid-sentence fragment.
        """
        import re
        # Split on sentence-ending punctuation: Chinese + English
        # Also split on 】 (Chinese doc section headers) and ； (Chinese semicolon)
        sentence_end = r'(?<=[.!?。！？\n；;】])\s*'
        raw_parts = re.split(sentence_end, text)
        # Merge tiny fragments back
        sentences = []
        buf = ""
        for part in raw_parts:
            if not part.strip():
                continue
            if buf and len(buf) < 20:  # short fragment — merge with previous
                buf += part
            else:
                if buf.strip():
                    sentences.append(buf.strip())
                buf = part
        if buf.strip():
            sentences.append(buf.strip())

        if not sentences:
            return [text]

        chunks = []
        current_sentences = []  # track sentence indices in current chunk
        current_len = 0

        for i, sent in enumerate(sentences):
            if current_len + len(sent) > self.chunk_size and current_sentences:
                # Finalize current chunk
                chunk_text = "".join(sentences[j] for j in current_sentences)
                chunks.append(chunk_text)

                # Overlap: carry over the last 1-2 complete sentences
                overlap_chars = 0
                overlap_sentences = []
                for j in reversed(current_sentences):
                    s = sentences[j]
                    if overlap_chars + len(s) <= self.chunk_overlap:
                        overlap_sentences.insert(0, j)
                        overlap_chars += len(s)
                    else:
                        break
                # Rebuild current from overlap sentences + new sentence
                current_sentences = list(overlap_sentences)
                current_len = sum(len(sentences[j]) for j in current_sentences)

            current_sentences.append(i)
            current_len += len(sent)

        if current_sentences:
            chunk_text = "".join(sentences[j] for j in current_sentences)
            chunks.append(chunk_text)

        return chunks if chunks else [text]